Update app.py

app.py

@@ -15,6 +15,10 @@ from PIL import Image
 from torchvision import transforms
 from torchvision.models import resnet50, ResNet50_Weights
 from torchvision import transforms, models
+from sentence_transformers import CrossEncoder
+
+
+reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
 
 class GeometryImageClassifier:
     def __init__(self):
@@ -155,6 +159,23 @@ def create_db(splits):
     vectordb = FAISS.from_documents(splits, embeddings)
     return vectordb
 
+def rerank_documents(query, docs, top_k=3):
+    pairs = [[query, doc.page_content] for doc in docs]
+    scores = reranker.predict(pairs)
+    doc_score_pairs = list(zip(docs, scores))
+    ranked_docs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
+    return [doc for doc, score in ranked_docs[:top_k]]
+
+def filter_relevant_chunks(query, chunks, embeddings, threshold=0.5):
+    query_embedding = embeddings.embed_query(query)
+    filtered_chunks = []
+    for chunk in chunks:
+        chunk_embedding = embeddings.embed_query(chunk.page_content)
+        similarity = cosine_similarity([query_embedding], [chunk_embedding])[0][0]
+        if similarity > threshold:
+            filtered_chunks.append(chunk)
+    return filtered_chunks
+
 def retrieve_documents(query, retriever, embeddings):
     print("\n=== Document Retrieval Process ===")
     print(f"Query: {query}")
@@ -208,7 +229,6 @@ def validate_query_semantically(query, retrieved_docs):
     
     return similarity_score >= 0.3
 
-
 def handle_query(query, history, retriever, qa_chain, embeddings):
     """ ✅ Handles user queries & prevents hallucination. """
     retrieved_docs = retrieve_documents(query, retriever, embeddings)