Create medical_rag.py

medical_rag.py

@@ -0,0 +1,159 @@
+from fastapi import HTTPException
+from pydantic import BaseModel
+import nltk
+from transformers import (
+    AutoTokenizer,
+    AutoModelForTokenClassification,
+    pipeline
+)
+from typing import List, Dict, Optional
+from general_rag import app, models, data, get_completion
+
+# Initialize NLTK
+nltk.download('punkt')
+
+class MedicalProfile(BaseModel):
+    conditions: str
+    daily_symptoms: str
+    count: int
+
+def load_medical_models():
+    try:
+        print("Loading medical domain models...")
+        
+        # Medical-specific models (only NER, no LLM)
+        models['bio_tokenizer'] = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
+        models['bio_model'] = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
+        models['ner_pipeline'] = pipeline("ner", model=models['bio_model'], tokenizer=models['bio_tokenizer'])
+        
+        print("Medical domain models loaded successfully")
+        return True
+    except Exception as e:
+        print(f"Error loading medical models: {e}")
+        return False
+
+def extract_entities(text):
+    try:
+        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)
+    except Exception as e:
+        print(f"Error extracting entities: {e}")
+        return []
+
+def match_entities(query_entities, sentence_entities):
+    try:
+        query_set, sentence_set = set(query_entities), set(sentence_entities)
+        matches = query_set.intersection(sentence_set)
+        return len(matches)
+    except Exception as e:
+        print(f"Error matching entities: {e}")
+        return 0
+
+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}")
+    
+    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)
+            if relevance_score >= min_query_words:
+                start_idx = max(0, i - portion_size // 2)
+                end_idx = min(len(sentences), i + portion_size // 2 + 1)
+                portion = " ".join(sentences[start_idx:end_idx])
+                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)
+            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 enhance_passage_with_entities(passage, entities):
+    return f"{passage}\n\nEntities: {', '.join(entities)}"
+
+def create_medical_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.
+
+    Passage: {passage}
+
+    Question: {question}
+
+    Answer:
+    """)
+    return prompt.format(passage=passage, question=question)
+
+@app.post("/api/chat")
+async def chat_endpoint(chat_query: ChatQuery):
+    try:
+        query_text = chat_query.query
+        language_code = chat_query.language_code        
+        if language_code == 0:
+            query_text = translate_text(query_text, 'ar_to_en')
+        
+        # Generate embeddings and retrieve relevant documents
+        query_embedding = embed_query_text(query_text)
+        n_results = 5
+        embeddings_data = load_embeddings()
+        folder_path = 'downloaded_articles/downloaded_articles'
+        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)
+        
+        # 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)
+        
+        # Extract relevant portions from documents using medical-specific function
+        relevant_portions = extract_relevant_portions(document_texts, query_text)
+        flattened_relevant_portions = []
+        for doc_id, portions in relevant_portions.items():
+            flattened_relevant_portions.extend(portions)
+        
+        combined_parts = " ".join(flattened_relevant_portions)
+        entities = extract_entities(query_text)
+        passage = enhance_passage_with_entities(combined_parts, entities)
+        
+        # Create medical-specific prompt and get completion from DeepSeek
+        prompt = create_medical_prompt(query_text, passage)
+        answer = get_completion(prompt)
+        
+        final_answer = answer.strip()
+        if language_code == 0:
+            final_answer = translate_text(final_answer, 'en_to_ar')
+            
+        if not final_answer:
+            final_answer = "Sorry, I can't help with that."
+            
+        return {
+            "response": f"I hope this answers your question: {final_answer}",
+            "success": True
+        }
+        
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Initialize medical models when this module is imported
+load_medical_models()