Update medical_rag.py

medical_rag.py

@@ -32,9 +32,10 @@ def load_medical_models():
         print(f"Error loading medical models: {e}")
         return False
 
-def extract_entities(text):
+def extract_entities(text, ner_pipeline=None):
     try:
-        ner_pipeline = models['ner_pipeline']
+        if ner_pipeline is None:
+            ner_pipeline = models['ner_pipeline']
         ner_results = ner_pipeline(text)
         entities = {result['word'] for result in ner_results if result['entity'].startswith("B-")}
         return list(entities)
@@ -55,13 +56,11 @@ def extract_relevant_portions(document_texts, query, max_portions=3, portion_siz
     relevant_portions = {}
     query_entities = extract_entities(query)
     print(f"Extracted Query Entities: {query_entities}")
-    
     for doc_id, doc_text in enumerate(document_texts):
         sentences = nltk.sent_tokenize(doc_text)  
         doc_relevant_portions = []
         doc_entities = extract_entities(doc_text)
         print(f"Document {doc_id} Entities: {doc_entities}")
-        
         for i, sentence in enumerate(sentences):
             sentence_entities = extract_entities(sentence)
             relevance_score = match_entities(query_entities, sentence_entities)
@@ -72,26 +71,47 @@ def extract_relevant_portions(document_texts, query, max_portions=3, portion_siz
                 doc_relevant_portions.append(portion)
             if len(doc_relevant_portions) >= max_portions:
                 break
-        
         if not doc_relevant_portions and len(doc_entities) > 0:
             print(f"Fallback: Selecting sentences with most entities for Document {doc_id}")
-            sorted_sentences = sorted(sentences, key=lambda s: len(extract_entities(s)), reverse=True)
+            sorted_sentences = sorted(sentences, key=lambda s: len(extract_entities(s, ner_biobert)), reverse=True)
             for fallback_sentence in sorted_sentences[:max_portions]:
                 doc_relevant_portions.append(fallback_sentence)
-        
         relevant_portions[f"Document_{doc_id}"] = doc_relevant_portions
     return relevant_portions
 
+def remove_duplicates(selected_parts):
+    unique_sentences = set()
+    unique_selected_parts = []
+    for sentence in selected_parts:
+        if sentence not in unique_sentences:
+            unique_selected_parts.append(sentence)
+            unique_sentences.add(sentence)
+    return unique_selected_parts
+
+def extract_entities(text):
+    try:
+        biobert_tokenizer = models['bio_tokenizer']
+        biobert_model = models['bio_model']
+        inputs = biobert_tokenizer(text, return_tensors="pt")
+        outputs = biobert_model(**inputs)
+        predictions = torch.argmax(outputs.logits, dim=2)
+        tokens = biobert_tokenizer.convert_ids_to_tokens(inputs.input_ids[0])
+        entities = [
+            tokens[i] 
+            for i in range(len(tokens)) 
+            if predictions[0][i].item() != 0  # Assuming 0 is the label for non-entity
+        ]
+        return entities
+    except Exception as e:
+        print(f"Error extracting entities: {e}")
+        return []
+
 def enhance_passage_with_entities(passage, entities):
     return f"{passage}\n\nEntities: {', '.join(entities)}"
 
-def create_medical_prompt(question, passage):
+def create_prompt(question, passage):
     prompt = ("""
-    As a medical expert, you are required to answer the following question based only on the provided passage. 
-    Do not include any information not present in the passage. 
-    Your response should directly reflect the content of the passage. 
-    Maintain accuracy and relevance to the provided information.
-    Provide a medically reliable answer in no more than 250 words.
+    As a medical expert, you are required to answer the following question based only on the provided passage. Do not include any information not present in the passage. Your response should directly reflect the content of the passage. Maintain accuracy and relevance to the provided information.
 
     Passage: {passage}
 
@@ -101,112 +121,184 @@ def create_medical_prompt(question, passage):
     """)
     return prompt.format(passage=passage, question=question)
 
+def generate_answer(prompt, max_length=860, temperature=0.2):
+    tokenizer_f = models['llm_tokenizer']
+    model_f = models['llm_model']
+    inputs = tokenizer_f(prompt, return_tensors="pt", truncation=True)
+    output_ids = model_f.generate(
+        inputs.input_ids,
+        max_length=max_length,
+        num_return_sequences=1,
+        temperature=temperature,
+        pad_token_id=tokenizer_f.eos_token_id
+    )
+    answer = tokenizer_f.decode(output_ids[0], skip_special_tokens=True)
+    passage_keywords = set(prompt.lower().split())  
+    answer_keywords = set(answer.lower().split())
+    if passage_keywords.intersection(answer_keywords):
+        return answer  
+    else:
+        return "Sorry, I can't help with that." 
+  
+def remove_answer_prefix(text):
+    prefix = "Answer:"
+    if prefix in text:
+        return text.split(prefix, 1)[-1].strip()  
+    return text
+
+def remove_incomplete_sentence(text):
+    if not text.endswith('.'):
+        last_period_index = text.rfind('.')
+        if last_period_index != -1:
+            return text[:last_period_index + 1].strip()
+    return text
+
 @app.post("/api/chat")
 async def chat_endpoint(chat_query: ChatQuery):
     try:
-        print("\n=== STARTING CHAT REQUEST PROCESSING ===")
-        print(f"Initial query: {chat_query.query} (language_code: {chat_query.language_code})")
-        
-        # Step 1: Handle translation if needed
         query_text = chat_query.query
-        language_code = chat_query.language_code
+        language_code = chat_query.language_code        
         if language_code == 0:
-            print("Translating from Arabic to English...")
-            query_text = translate_text(query_text, 'ar_to_en')
-            print(f"Translated query: {query_text}")
+            query_text = translate_ar_to_en(query_text)
         
-        # Step 2: Generate embeddings
-        print("\nGenerating query embeddings...")
+        # Generate embeddings and retrieve relevant documents (original RAG logic)
         query_embedding = embed_query_text(query_text)
-        print(f"Embedding generated. Shape: {query_embedding.shape}")
-        
-        # Step 3: Load embeddings and query them
-        print("\nLoading document embeddings...")
-        embeddings_data = load_embeddings()
-        if not embeddings_data:
-            raise HTTPException(status_code=500, detail="Failed to load embeddings data")
-        print(f"Loaded embeddings for {len(embeddings_data)} documents")
-        
-        print("\nQuerying embeddings...")
         n_results = 5
+        embeddings_data = load_embeddings()
+        folder_path = 'downloaded_articles/downloaded_articles'
         initial_results = query_embeddings(query_embedding, embeddings_data, n_results)
-        print(f"Initial results: {initial_results}")
-        
         document_ids = [doc_id for doc_id, _ in initial_results]
-        print(f"Document IDs to retrieve: {document_ids}")
-        
-        # Step 4: Retrieve document texts
-        print("\nRetrieving document texts...")
-        folder_path = 'downloaded_articles/downloaded_articles'
         document_texts = retrieve_document_texts(document_ids, folder_path)
-        print(f"Retrieved {len(document_texts)} documents")
-        
-        # Step 5: Rerank documents
-        print("\nReranking documents...")
-        if 'cross_encoder' not in models:
-            raise HTTPException(status_code=500, detail="Cross-encoder model not loaded")
         
+        # Rerank documents with cross-encoder
         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)
-        print("Top 3 reranked documents:")
-        for i, (score, doc_id, _) in enumerate(scored_documents[:3]):
-            print(f"{i+1}. Doc {doc_id} (score: {score:.4f})")
-        
-        # Step 6: Extract relevant portions
-        print("\nExtracting relevant portions...")
-        relevant_portions = extract_relevant_portions(document_texts, query_text)
-        print(f"Found relevant portions in {len(relevant_portions)} documents")
         
+        # Extract relevant portions from documents
+        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)
+        unique_selected_parts = remove_duplicates(flattened_relevant_portions)
+        combined_parts = " ".join(unique_selected_parts)
         
-        combined_parts = " ".join(flattened_relevant_portions)
-        print(f"Combined relevant text length: {len(combined_parts)} characters")
-        
-        # Step 7: Extract and enhance with entities
-        print("\nExtracting entities...")
+        # Enhance context with entities
         entities = extract_entities(query_text)
-        print(f"Found entities: {entities}")
-        
         passage = enhance_passage_with_entities(combined_parts, entities)
-        print(f"Enhanced passage length: {len(passage)} characters")
         
-        # Step 8: Generate response
-        print("\nCreating prompt...")
-        prompt = create_medical_prompt(query_text, passage)
-        print(f"Prompt length: {len(prompt)} characters")
+        # Create prompt and generate answer using OpenRouter
+        prompt = create_prompt(query_text, passage)
         
-        print("\nGetting completion from DeepSeek...")
-        answer = get_completion(prompt)
-        print(f"Raw answer received: {answer[:200]}...")  # Print first 200 chars
+        # Add constraints similar to /api/ask endpoint
+        constraints = "Provide a medically reliable answer in no more than 250 words."
+        full_prompt = f"{prompt} {constraints}"
         
-        # Step 9: Final processing
+        # Use the same OpenRouter model as /api/ask
+        answer = get_completion(full_prompt)
+        
+        # Process the answer
         final_answer = answer.strip()
         if language_code == 0:
-            print("\nTranslating answer to Arabic...")
-            final_answer = translate_text(final_answer, 'en_to_ar')
-            print(f"Translated answer: {final_answer[:200]}...")
+            final_answer = translate_en_to_ar(final_answer)
             
         if not final_answer:
             final_answer = "Sorry, I can't help with that."
-            print("Warning: Empty answer received")
-        
-        print("\n=== REQUEST PROCESSING COMPLETE ===")
+            
         return {
             "response": f"I hope this answers your question: {final_answer}",
             "success": True
         }
         
     except HTTPException as e:
-        print(f"\n!!! HTTPException: {e.detail}")
         raise e
     except Exception as e:
-        print(f"\n!!! Unexpected error: {str(e)}")
-        print(f"Error type: {type(e).__name__}")
         raise HTTPException(status_code=500, detail=str(e))
+        
+
+@app.post("/api/resources")
+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)
+        if not initial_results:
+            raise ValueError("No relevant documents found.")
+        document_ids = [doc_id for doc_id, _ in initial_results]
+        file_path = 'finalcleaned_excel_file.xlsx'
+        df = pd.read_excel(file_path)
+        file_name_to_url = {f"article_{index}.html": url for index, url in enumerate(df['Unnamed: 0'])}
+        resources = []
+        for file_name in document_ids:
+            original_url = file_name_to_url.get(file_name, None)
+            if original_url:
+                title = get_page_title(original_url) or "Unknown Title"
+                resources.append({"file_name": file_name, "title": title, "url": original_url})
+            else:
+                resources.append({"file_name": file_name, "title": "Unknown", "url": None})        
+        document_texts = retrieve_document_texts(document_ids, 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, resource in enumerate(resources):
+            resource["score"] = scores[i] if i < len(scores) else 0.0        
+        resources.sort(key=lambda x: x["score"], reverse=True)       
+        output = [{"title": resource["title"], "url": resource["url"]} for resource in resources]       
+        return output
+    except ValueError as ve:
+        raise HTTPException(status_code=400, detail=str(ve))
+    except Exception as e:
+        print(f"Unexpected error: {e}")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
+
+@app.post("/api/recipes")
+async def recipes_endpoint(profile: MedicalProfile):
+    try:
+        recipe_query = (
+            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)
+        if not initial_results:
+            raise ValueError("No relevant recipes found.")
+        print("Initial results (document indices and similarities):")
+        print(initial_results)
+        document_indices = [doc_id for doc_id, _ in initial_results]
+        print("Document indices:", document_indices)
+        metadata_path = 'recipes_metadata.xlsx'
+        metadata = retrieve_metadata(document_indices, metadata_path=metadata_path)
+        print(f"Retrieved Metadata: {metadata}")        
+        recipes = []
+        for item in metadata.values():
+            recipes.append({
+                "title": item["original_file_name"] if "original_file_name" in item else "Unknown Title",
+                "url": item["url"] if "url" in item else ""
+            })               
+        print(recipes)
+        return recipes
+    except ValueError as ve:
+        raise HTTPException(status_code=400, detail=str(ve))
+    except Exception as e:
+        print(f"Unexpected error: {e}")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
+      
 
 # Initialize medical models when this module is imported
 load_medical_models()