Update app.py

app.py

@@ -350,47 +350,41 @@ from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
 import nltk
 
-# Load a pre-trained embedding model
-embedding_model = SentenceTransformer('all-MiniLM-L6-v2')  # Use a lightweight model for speed
-
-from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
-
-# Load the embedding model globally for efficiency
-embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
-
-def extract_relevant_portions(document_texts, query, max_portions=3, chunk_size=500):
+def extract_relevant_portions(query_embedding, top_documents, embeddings_data, max_portions=3):
     try:
-        # Embed the query once
-        query_embedding = embedding_model.encode([query])
-
         relevant_portions = {}
-        for doc_id, doc_text in enumerate(document_texts):
-            # Split document into chunks (e.g., 500 characters per chunk)
-            chunks = [doc_text[i:i + chunk_size] for i in range(0, len(doc_text), chunk_size)]
-            
-            # Embed all chunks in a single batch
-            chunk_embeddings = embedding_model.encode(chunks)
+
+        for _, doc_id, doc_text in top_documents:
+            if doc_id not in embeddings_data:
+                print(f"Warning: No embedding available for Document ID {doc_id}. Skipping...")
+                continue
             
-            # Compute cosine similarity between query and all chunks
-            similarities = cosine_similarity(query_embedding, chunk_embeddings)[0]
+            # Retrieve the precomputed embedding for this document
+            doc_embedding = np.array(embeddings_data[doc_id])
             
-            # Rank chunks by similarity
-            ranked_chunks = sorted(
-                enumerate(chunks),
-                key=lambda x: similarities[x[0]],
-                reverse=True
-            )
+            # Compute similarity between the query embedding and the document embedding
+            similarity = cosine_similarity(query_embedding, [doc_embedding]).flatten()[0]
+
+            # Split the document into sentences
+            sentences = nltk.sent_tokenize(doc_text)
+
+            # Rank sentences based on their length (proxy for importance) or other heuristic
+            # Since we're using document-level embeddings, we assume all sentences are equally relevant.
+            sorted_sentences = sorted(sentences, key=lambda x: len(x), reverse=True)[:max_portions]
+
+            relevant_portions[doc_id] = sorted_sentences
             
-            # Select top chunks based on similarity
-            doc_relevant_portions = [chunk for _, chunk in ranked_chunks[:max_portions]]
-            relevant_portions[f"Document_{doc_id}"] = doc_relevant_portions
+            print(f"Extracted relevant portions for Document ID {doc_id} (Similarity: {similarity:.4f}):")
+            for i, sentence in enumerate(sorted_sentences, start=1):
+                print(f"  Portion {i}: {sentence[:100]}...")  # Print first 100 characters for preview
 
         return relevant_portions
+
     except Exception as e:
-        print(f"Error in extracting relevant portions: {e}")
+        print(f"Error in extract_relevant_portions: {e}")
         return {}
 
+
 def remove_duplicates(selected_parts):
     unique_sentences = set()
     unique_selected_parts = []
@@ -532,7 +526,7 @@ async def chat_endpoint(chat_query: ChatQuery):
         scores = cross_encoder.predict([(query_text, doc) for doc in document_texts])
         scored_documents = list(zip(scores, document_ids, document_texts))
         scored_documents.sort(key=lambda x: x[0], reverse=True)
-        relevant_portions = extract_relevant_portions(document_texts, query_text, max_portions=3, chunk_size=500)
+        relevant_portions = extract_relevant_portions(query_embedding, scored_documents, embeddings_data, max_portions=3)
         #flattened_relevant_portions = []
         #for doc_id, portions in relevant_portions.items():
             #flattened_relevant_portions.extend(portions)