Update app.py

app.py

@@ -3,8 +3,6 @@ import chromadb
 import torch
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.retrievers import ContextualCompressionRetriever
-from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 import gradio as gr
 
@@ -31,11 +29,6 @@ reranker = HuggingFaceCrossEncoder(
     model_name="BAAI/bge-reranker-base",
     model_kwargs={"device": device}
 )
-compressor = CrossEncoderReranker(model=reranker, top_n=5)
-retriever = vectorstore.as_retriever(search_kwargs={"k": 10})  # Retrieve 2k initially
-compression_retriever = ContextualCompressionRetriever(
-    base_compressor=compressor, base_retriever=retriever
-)
 
 def add_text_to_db(text):
     """
@@ -84,20 +77,30 @@ def search_similar_texts(query, k, threshold):
     if not isinstance(threshold, (int, float)) or threshold < 0 or threshold > 1:
         return "Error: Threshold must be a number between 0 and 1."
     
-    # Retrieve and rerank
-    retriever.search_kwargs["k"] = max(k * 2, 10)  # Retrieve 2k or at least 10
-    compressor.top_n = k  # Rerank to top k
-    docs = compression_retriever.get_relevant_documents(query)
+    # Retrieve initial documents
+    retriever = vectorstore.as_retriever(search_kwargs={"k": max(k * 2, 10)})
+    docs = retriever.get_relevant_documents(query)
     
     if not docs:
         return "No such record."
     
-    # Filter results by threshold
-    results = []
-    for i, doc in enumerate(docs[:k]):  # Ensure at most k results
+    # Compute reranker scores
+    scored_docs = []
+    for doc in docs:
         text = doc.metadata.get("text", "No text available")
-        score = doc.metadata.get("score", 0.0)  # Reranker score
+        # Prepare input for reranker: list of [query, document] pairs
+        score = reranker.predict([[query, text]])[0]
+        doc.metadata["score"] = float(score)
+        scored_docs.append((doc, score))
+    
+    # Sort by score in descending order
+    scored_docs.sort(key=lambda x: x[1], reverse=True)
+    
+    # Filter by threshold and limit to k
+    results = []
+    for i, (doc, score) in enumerate(scored_docs[:k]):
         if score >= threshold:
+            text = doc.metadata.get("text", "No text available")
             results.append(f"Result {i+1}:\nText: {text}\nScore: {score:.4f}\n")
     
     if not results: