medical_rag.py

from fastapi import HTTPException
from pydantic import BaseModel
import torch 
import nltk
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    pipeline
)
from typing import List, Dict, Optional
from general_rag import app, models, data, ChatQuery, embed_query_text, get_completion, load_embeddings, translate_text, retrieve_document_texts, query_embeddings

# 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, ner_pipeline=None):
    try:
        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)
    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, 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_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.

    Passage: {passage}

    Question: {question}

    Answer:
    """)
    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:
        query_text = chat_query.query
        language_code = chat_query.language_code        
        if language_code == 0:
            query_text = translate_ar_to_en(query_text)
        
        # Generate embeddings and retrieve relevant documents (original RAG logic)
        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
        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)
        
        # Enhance context with entities
        entities = extract_entities(query_text)
        passage = enhance_passage_with_entities(combined_parts, entities)
        
        # Create prompt and generate answer using OpenRouter
        prompt = create_prompt(query_text, passage)
        
        # Add constraints similar to /api/ask endpoint
        constraints = "Provide a medically reliable answer in no more than 250 words."
        full_prompt = f"{prompt} {constraints}"
        
        # Use the same OpenRouter model as /api/ask
        answer = get_completion(full_prompt)
        
        # Process the answer
        final_answer = answer.strip()
        if language_code == 0:
            final_answer = translate_en_to_ar(final_answer)
            
        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))
        

@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()