Update app.py

app.py

@@ -51,11 +51,12 @@ class QueryRequest(BaseModel):
 class MedicalProfile(BaseModel):
     conditions: str
     daily_symptoms: str
+    count: int
 
 class ChatQuery(BaseModel):
     query: str
     language_code: int = 1
-    conversation_id: str
+    #conversation_id: str
 
 class ChatMessage(BaseModel):
     role: str
@@ -219,7 +220,7 @@ def embed_query_text(query_text):
     query_embedding = embedding.encode([query_text])
     return query_embedding
 
-def query_embeddings(query_embedding, embeddings_data=None, n_results=5):
+def query_embeddings(query_embedding, embeddings_data=None, n_results):
     embeddings_data = load_embeddings()
     if not embeddings_data:
         print("No embeddings data available.")
@@ -234,7 +235,7 @@ def query_embeddings(query_embedding, embeddings_data=None, n_results=5):
         print(f"Error in query_embeddings: {e}")
         return []
 
-def query_recipes_embeddings(query_embedding, embeddings_data, n_results = 5):
+def query_recipes_embeddings(query_embedding, embeddings_data, n_results):
     embeddings_data = load_recipes_embeddings()
     if embeddings_data is None:
         print("No embeddings data available.")
@@ -365,7 +366,7 @@ def match_entities(query_entities, sentence_entities):
         print(f"Error matching entities: {e}")
         return 0
 
-def extract_relevant_portions(document_texts, query, max_portions=3, portion_size=1, min_query_words=1):
+def extract_relevant_portions(document_texts, query, max_portions=3, portion_size=1, min_query_words=2):
     relevant_portions = {}
     query_entities = extract_entities(query)
     print(f"Extracted Query Entities: {query_entities}")
@@ -466,6 +467,40 @@ def remove_incomplete_sentence(text):
             return text[:last_period_index + 1].strip()
     return text
 
+def translate_ar_to_en(text):
+    try:
+        ar_to_en_tokenizer = models['ar_to_en_tokenizer'] = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
+        ar_to_en_model= models['ar_to_en_model'] = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
+        inputs = ar_to_en_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+        translated_ids = ar_to_en_model.generate(
+            inputs.input_ids,
+            max_length=512, 
+            num_beams=4,     
+            early_stopping=True
+        )
+        translated_text = ar_to_en_tokenizer.decode(translated_ids[0], skip_special_tokens=True)
+        return translated_text
+    except Exception as e:
+        print(f"Error during Arabic to English translation: {e}")
+        return None
+              
+def translate_en_to_ar(text):
+    try:
+        en_to_ar_tokenizer = models['en_to_ar_tokenizer'] = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
+        en_to_ar_model = models['en_to_ar_model'] = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ar")  
+        inputs = en_to_ar_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+        translated_ids = en_to_ar_model.generate(
+            inputs.input_ids,
+            max_length=512,  
+            num_beams=4,     
+            early_stopping=True
+        )
+        translated_text = en_to_ar_tokenizer.decode(translated_ids[0], skip_special_tokens=True)
+        return translated_text
+    except Exception as e:
+        print(f"Error during English to Arabic translation: {e}")
+        return None
+
 @app.get("/")
 async def root():
     return {"message": "Welcome to the FastAPI application! Use the /health endpoint to check health, and /api/query for processing queries."}
@@ -485,18 +520,20 @@ async def health_check():
 async def chat_endpoint(chat_query: ChatQuery):
     try:
         query_text = chat_query.query
-        language_code = chat_query.language_code
+        language_code = chat_query.language_code        
+        if language_code == 0:
+            query_text = translate_ar_to_en(query_text)
         query_embedding = embed_query_text(query_text)  
         embeddings_data = load_embeddings ()
         folder_path = 'downloaded_articles/downloaded_articles'
-        initial_results = query_embeddings(query_embedding, embeddings_data, n_results=5)
+        initial_results = query_embeddings(query_embedding, embeddings_data, n_results)
         document_ids = [doc_id for doc_id, _ in initial_results]
         document_texts = retrieve_document_texts(document_ids, folder_path)
         cross_encoder = models['cross_encoder']
         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, portion_size=1, min_query_words=1)
+        relevant_portions = extract_relevant_portions(document_texts, query_text, max_portions=3, portion_size=1, min_query_words=2)
         flattened_relevant_portions = []
         for doc_id, portions in relevant_portions.items():
             flattened_relevant_portions.extend(portions)
@@ -518,8 +555,8 @@ async def chat_endpoint(chat_query: ChatQuery):
         else:
             print("Sorry, I can't help with that.")
         return {
-            "response": final_answer,
-            "conversation_id": chat_query.conversation_id,
+            "response": "I hope this answers your question:" {final_answer},
+            #"conversation_id": chat_query.conversation_id,
             "success": True
         }
     except Exception as e:
@@ -529,13 +566,14 @@ async def chat_endpoint(chat_query: ChatQuery):
 async def resources_endpoint(profile: MedicalProfile):
     try:       
         query_text = profile.conditions + " " + profile.daily_symptoms
+        n_results = profile.count
         print(f"Generated query text: {query_text}")        
         query_embedding = embed_query_text(query_text)
         if query_embedding is None:
             raise ValueError("Failed to generate query embedding.")
         embeddings_data = load_embeddings()
         folder_path = 'downloaded_articles/downloaded_articles'
-        initial_results = query_embeddings(query_embedding, embeddings_data, n_results=6)
+        initial_results = query_embeddings(query_embedding, embeddings_data, n_results)
         if not initial_results:
             raise ValueError("No relevant documents found.")
         document_ids = [doc_id for doc_id, _ in initial_results]
@@ -570,17 +608,18 @@ async def resources_endpoint(profile: MedicalProfile):
 async def recipes_endpoint(profile: MedicalProfile):
     try:
         recipe_query = (
-            f"Recipes foods and meals suitable for someone with: "
+            f"Recipes and foods for: "
             f"{profile.conditions} and experiencing {profile.daily_symptoms}"
         )
         query_text = recipe_query
         print(f"Generated query text: {query_text}")
+        n_results = profile.count
         query_embedding = embed_query_text(query_text)
         if query_embedding is None:
             raise ValueError("Failed to generate query embedding.")
         embeddings_data = load_recipes_embeddings()
         folder_path = 'downloaded_articles/downloaded_articles'
-        initial_results = query_recipes_embeddings(query_embedding, embeddings_data, n_results=5)
+        initial_results = query_recipes_embeddings(query_embedding, embeddings_data, n_results)
         if not initial_results:
             raise ValueError("No relevant recipes found.")
         print("Initial results (document indices and similarities):")
@@ -590,8 +629,19 @@ async def recipes_endpoint(profile: MedicalProfile):
         metadata_path = 'recipes_metadata.xlsx'
         metadata = retrieve_metadata(document_indices, metadata_path=metadata_path)
         print(f"Retrieved Metadata: {metadata}")
+        recipes = metadata
+        document_texts = retrieve_rec_texts(document_indices, folder_path)
+        if not document_texts:
+            raise ValueError("Failed to retrieve document texts.")        
+        cross_encoder = models['cross_encoder']
+        scores = cross_encoder.predict([(query_text, doc) for doc in document_texts])
+        scores = [float(score) for score in scores]       
+        for i, recipe in enumerate(recipes):
+            recipe["score"] = scores[i] if i < len(scores) else 0.0
+        recipes.sort(key=lambda x: x["score"], reverse=True)       
+        return {"recipes": recipes[:5], "success": True}
         response = {           
-            "metadata": metadata,
+            "recipes": recipes,
         }
         return response       
     except ValueError as ve: