Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,10 +3,12 @@ import re
 import random
 import PyPDF2
 import numpy as np
-from collections import defaultdict, deque
-from sentence_transformers import SentenceTransformer
+from collections import defaultdict
 from sklearn.metrics.pairwise import cosine_similarity
 
+import torch
+from transformers import AutoTokenizer, AutoModel
+
 # ---------------------
 # Tokenization
 # ---------------------
@@ -20,13 +22,15 @@ class PDFQASystem:
     def __init__(self):
         self.text_chunks = []
         self.embeddings = None
-        self.model = SentenceTransformer('all-MiniLM-L6-v2')
+        self.tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+        self.model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+        self.model.eval()
         self.active_document = None
-        
+
     def process_pdf_stream(self, uploaded_file):
         text = self._extract_pdf_text(uploaded_file)
         self.text_chunks = self._chunk_text(text)
-        self.embeddings = self.model.encode(self.text_chunks)
+        self.embeddings = self._embed(self.text_chunks)
         self.active_document = uploaded_file.name
 
     def _extract_pdf_text(self, uploaded_file):
@@ -42,11 +46,23 @@ class PDFQASystem:
         words = text.split()
         return [' '.join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
 
+    def _mean_pooling(self, model_output, attention_mask):
+        token_embeddings = model_output.last_hidden_state
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        return torch.sum(token_embeddings * input_mask_expanded, dim=1) / torch.clamp(input_mask_expanded.sum(dim=1), min=1e-9)
+
+    def _embed(self, texts):
+        inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
+        with torch.no_grad():
+            model_output = self.model(**inputs)
+        embeddings = self._mean_pooling(model_output, inputs['attention_mask'])
+        return torch.nn.functional.normalize(embeddings, p=2, dim=1).cpu().numpy()
+
     def answer_question(self, question):
         if not self.active_document:
             return "No document loaded. Please upload a PDF first."
-        
-        question_embedding = self.model.encode(question)
+
+        question_embedding = self._embed([question])[0]
         similarities = cosine_similarity([question_embedding], self.embeddings)[0]
         best_match_idx = np.argmax(similarities)
         return self.text_chunks[best_match_idx]