notabaka commited on
Commit
79ecc72
1 Parent(s): 0c7ffdb
Files changed (2) hide show
  1. app.py +37 -27
  2. requirements.txt +4 -1
app.py CHANGED
@@ -3,6 +3,9 @@ import torch
3
  import torch.nn.functional as F
4
  from torch import Tensor
5
  from transformers import AutoTokenizer, AutoModel
 
 
 
6
 
7
  def last_token_pool(last_hidden_states: Tensor,
8
  attention_mask: Tensor) -> Tensor:
@@ -21,37 +24,44 @@ st.title("Text Similarity Model")
21
 
22
  task = 'Given a web search query, retrieve relevant passages that answer the query'
23
 
24
- query1 = st.text_input("Enter first query")
25
- query2 = st.text_input("Enter second query")
 
26
 
27
- if query1 and query2:
28
- queries = [
29
- get_detailed_instruct(task, query1),
30
- get_detailed_instruct(task, query2)
31
- ]
32
 
33
- passages = [
34
- "To bake a delicious chocolate cake, you'll need the following ingredients: all-purpose flour, sugar, cocoa powder, baking powder, baking soda, salt, eggs, milk, vegetable oil, and vanilla extract. Start by preheating your oven to 350°F (175°C). In a mixing bowl, combine the dry ingredients (flour, sugar, cocoa powder, baking powder, baking soda, and salt). In a separate bowl, whisk together the wet ingredients (eggs, milk, vegetable oil, and vanilla extract). Gradually add the wet mixture to the dry ingredients, stirring until well combined. Pour the batter into a greased cake pan and bake for 30-35 minutes. Let it cool before frosting with your favorite chocolate frosting. Enjoy your homemade chocolate cake!",
35
- "The flu, or influenza, is an illness caused by influenza viruses. Common symptoms of the flu include a high fever, chills, cough, sore throat, runny or stuffy nose, body aches, headache, fatigue, and sometimes nausea and vomiting. These symptoms can come on suddenly and are usually more severe than the common cold. It's important to get plenty of rest, stay hydrated, and consult a healthcare professional if you suspect you have the flu. In some cases, antiviral medications can help alleviate symptoms and reduce the duration of the illness."]
 
36
 
37
- tokenizer = AutoTokenizer.from_pretrained('Salesforce/SFR-Embedding-Mistral')
38
- model = AutoModel.from_pretrained('Salesforce/SFR-Embedding-Mistral')
39
-
40
- # Get embeddings
41
- max_length = 4096
42
- input_texts = queries + passages
43
- batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors="pt")
44
-
45
- outputs = model(**batch_dict)
46
- embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
47
-
48
- # Normalize embeddings
49
- embeddings = F.normalize(embeddings, p=2, dim=1)
50
-
51
- scores = (embeddings[:2] @ embeddings[2:].T) * 100
52
-
53
- st.write("Similarity scores:", scores.tolist())
54
 
 
 
55
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56
 
57
 
 
3
  import torch.nn.functional as F
4
  from torch import Tensor
5
  from transformers import AutoTokenizer, AutoModel
6
+ import textract
7
+ import docx2txt
8
+ import pdfplumber
9
 
10
  def last_token_pool(last_hidden_states: Tensor,
11
  attention_mask: Tensor) -> Tensor:
 
24
 
25
  task = 'Given a web search query, retrieve relevant passages that answer the query'
26
 
27
+ docs = st.sidebar.file_uploader("Upload documents", accept_multiple_files=True, type=['txt','pdf','xlsx','docx'])
28
+ query = st.text_input("Enter search query")
29
+ click = st.button("Search")
30
 
31
+ if click and query:
32
+ doc_contents = []
 
 
 
33
 
34
+ for doc in docs:
35
+ # Extract text from each document
36
+ doc_text = extract_text(doc)
37
+ doc_contents.append(doc_text)
38
 
39
+ doc_embeddings = get_embeddings(doc_contents)
40
+ query_embedding = get_embedding(query)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
 
42
+ scores = compute_similarity(query_embedding, doc_embeddings)
43
+ ranked_docs = get_ranked_docs(scores)
44
 
45
+ st.write("Most Relevant Documents")
46
+ for doc, score in ranked_docs:
47
+ st.write(f"{doc.name} (score: {score:.2f})")
48
+
49
+ def extract_text(doc):
50
+ if doc.type == 'text/plain':
51
+ return doc.getvalue().decode("utf-8")
52
+
53
+ if doc.type == "application/pdf":
54
+ with pdfplumber.open(doc) as pdf:
55
+ pages = [page.extract_text() for page in pdf.pages]
56
+ return "\n".join(pages)
57
+
58
+ if doc.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
59
+ return docx2txt.process(doc)
60
+
61
+ if doc.name.endswith(".xlsx"):
62
+ text = textract.process(doc)
63
+ return text.decode("utf-8")
64
+
65
+ return None
66
 
67
 
requirements.txt CHANGED
@@ -1,2 +1,5 @@
1
  torch
2
- transformers
 
 
 
 
1
  torch
2
+ transformers
3
+ textract
4
+ docx2txt
5
+ pdfplumber