Update app.py

app.py

@@ -39,8 +39,6 @@ from typing import List, Dict, Tuple
 import datetime
 from abc import ABC, abstractmethod
 from typing import List, Dict, Any
-from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
-
 
 # Automatically get the current year
 CURRENT_YEAR = datetime.datetime.now().year
@@ -107,9 +105,8 @@ groq_client = Groq(api_key=GROQ_API_KEY)
 MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
 mistral_client = Mistral(api_key=MISTRAL_API_KEY)
 
-similarity_model = HuggingFaceInferenceAPIEmbeddings(
-    api_key=HF_TOKEN, model_name="sentence-transformers/all-MiniLM-l6-v2"
-)
+# Initialize the similarity model
+similarity_model = SentenceTransformer('BAAI/bge-small-en-v1.5')
 
 # Step 1: Create a base class for AI models
 class AIModel(ABC):
@@ -648,9 +645,9 @@ def rerank_documents(query: str, documents: List[Dict],
         bm25_scores = bm25.get_scores(query)
         
         # Step 4: Get semantic similarity scores
-        query_embedding = similarity_model.embed_query(query)
+        query_embedding = similarity_model.encode(query)
         doc_summaries = [doc['summary'] for doc in valid_docs]
-        doc_embeddings = similarity_model.embed_query(doc_summaries)
+        doc_embeddings = similarity_model.encode(doc_summaries)
         semantic_scores = util.cos_sim(query_embedding, doc_embeddings)[0]
         
         # Step 5: Combine scores (normalize first)
@@ -682,11 +679,11 @@ def rerank_documents(query: str, documents: List[Dict],
                 continue
                 
             # Check similarity with already selected documents
-            doc_embedding = similarity_model.embed_query(doc['summary'])
+            doc_embedding = similarity_model.encode(doc['summary'])
             is_similar = False
             
             for content in added_contents:
-                content_embedding = similarity_model.embed_query(content)
+                content_embedding = similarity_model.encode(content)
                 similarity = util.pytorch_cos_sim(doc_embedding, content_embedding)
                 if similarity > similarity_threshold:
                     is_similar = True
@@ -708,8 +705,8 @@ def rerank_documents(query: str, documents: List[Dict],
 
 def compute_similarity(text1, text2):
     # Encode the texts
-    embedding1 = similarity_model.embed_query(text1)
-    embedding2 = similarity_model.embed_query(text2)
+    embedding1 = similarity_model.encode(text1)
+    embedding2 = similarity_model.encode(text2)
     
     # Compute cosine similarity
     cosine_similarity = util.pytorch_cos_sim(embedding1, embedding2)