Update app.py

app.py

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