tststml
Browse files- app.py +37 -27
- 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 |
-
|
25 |
-
|
|
|
26 |
|
27 |
-
if
|
28 |
-
|
29 |
-
get_detailed_instruct(task, query1),
|
30 |
-
get_detailed_instruct(task, query2)
|
31 |
-
]
|
32 |
|
33 |
-
|
34 |
-
|
35 |
-
|
|
|
36 |
|
37 |
-
|
38 |
-
|
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
|