Delete model.py

model.py

@@ -1,323 +0,0 @@
-"""
-Model loading and RAG pipeline core functions.
-Handles vector store, embeddings, and answer generation.
-"""
-
-import os
-import re
-import torch
-import traceback
-import random
-import numpy as np
-from pathlib import Path
-
-try:
-    import chromadb
-except ImportError as e:
-    print(f"Warning: chromadb import failed: {e}")
-    chromadb = None
-
-try:
-    from sentence_transformers import SentenceTransformer
-except ImportError as e:
-    print(f"Error: sentence_transformers not available: {e}")
-    raise
-
-try:
-    from transformers import (
-        AutoTokenizer,
-        AutoModelForSeq2SeqLM,
-        pipeline
-    )
-except ImportError as e:
-    print(f"Error: transformers not available: {e}")
-    raise
-
-
-# ===========================
-# CONFIGURATION & SETUP
-# ===========================
-
-# Set random seeds for reproducibility
-random.seed(1)
-np.random.seed(1)
-torch.manual_seed(1)
-if torch.cuda.is_available():
-    torch.cuda.manual_seed_all(1)
-
-# Paths (adjust these to match your HuggingFace Hub paths)
-VECTOR_DB_PATH = "./MedQuAD_db"
-FINE_TUNED_MODEL_ID = "amiraghhh/fine-tuned-flan-t5-small"  # Update with your HF model path
-
-# Global objects (loaded once at startup)
-embed_model = None
-vector_store = None
-flant5_tokenizer = None
-flant5_model = None
-finetuned_llm = None
-rerank_tokenizer = None
-rerank_model = None
-rewriter_llm = None
-
-
-# ===========================
-# INITIALIZATION FUNCTIONS
-# ===========================
-
-def load_embeddings_model():
-    """Load SentenceTransformer embedding model."""
-    global embed_model
-    if embed_model is None:
-        print("Loading embedding model...")
-        embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-    return embed_model
-
-
-def load_vector_store():
-    """Load ChromaDB vector store from disk."""
-    global vector_store
-    if vector_store is None:
-        print("Loading vector store...")
-        db_client = chromadb.PersistentClient(path=VECTOR_DB_PATH)
-        try:
-            vector_store = db_client.get_collection("medical_rag")
-            print(f"Vector store loaded with {vector_store.count()} documents")
-        except Exception as e:
-            print(f"Error loading vector store: {e}")
-            raise
-    return vector_store
-
-
-def load_flan_t5_models():
-    """Load baseline FLAN-T5 models for prompt building."""
-    global flant5_tokenizer, flant5_model
-    if flant5_tokenizer is None:
-        print("Loading FLAN-T5 models...")
-        flant5_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
-        flant5_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
-    return flant5_tokenizer, flant5_model
-
-
-def load_rewriter_model():
-    """Load query rewriter model (FLAN-T5 small)."""
-    global rewriter_llm
-    if rewriter_llm is None:
-        print("Loading query rewriter...")
-        rewriter_llm = pipeline(
-            "text2text-generation",
-            model="google/flan-t5-small",
-            max_length=64,
-            do_sample=False,
-            temperature=0.3,
-            repetition_penalty=1.3,
-            no_repeat_ngram_size=2
-        )
-    return rewriter_llm
-
-
-def load_reranker_model():
-    """Load MonoT5 reranker model."""
-    global rerank_tokenizer, rerank_model
-    if rerank_tokenizer is None:
-        print("Loading reranker model...")
-        rerank_tokenizer = AutoTokenizer.from_pretrained("castorini/monot5-base-msmarco")
-        rerank_model = AutoModelForSeq2SeqLM.from_pretrained("castorini/monot5-base-msmarco")
-        rerank_model.eval()
-    return rerank_tokenizer, rerank_model
-
-
-def load_finetuned_model():
-    """Load fine-tuned FLAN-T5 model for answer generation."""
-    global finetuned_llm
-    if finetuned_llm is None:
-        print("Loading fine-tuned model...")
-        ft_tokenizer = AutoTokenizer.from_pretrained(FINE_TUNED_MODEL_ID)
-        ft_model = AutoModelForSeq2SeqLM.from_pretrained(FINE_TUNED_MODEL_ID)
-        
-        finetuned_llm = pipeline(
-            "text2text-generation",
-            model=ft_model,
-            tokenizer=ft_tokenizer,
-            decoder_start_token_id=ft_model.config.pad_token_id
-        )
-    return finetuned_llm
-
-
-def initialize_all():
-    """Load all models and vector store at startup."""
-    print("Initializing RAG pipeline...")
-    load_embeddings_model()
-    load_vector_store()
-    load_flan_t5_models()
-    load_rewriter_model()
-    load_reranker_model()
-    load_finetuned_model()
-    print("RAG pipeline initialized successfully!")
-
-
-# ===========================
-# PROMPT BUILDING
-# ===========================
-
-def build_prompt(user_query, context, rewritten_query, max_tokens=512):
-    """Build prompt with context and query within token limit.
-    
-    Args:
-        user_query (str): Original user question
-        context (list): Retrieved context chunks
-        rewritten_query (str): Query after rewriting
-        max_tokens (int): Maximum tokens for full prompt
-        
-    Returns:
-        str: Formatted prompt for the model
-    """
-    tokenizer, _ = load_flan_t5_models()
-    
-    if not context:
-        return f"""No relevant medical information found.
-Q: {rewritten_query}
-A: Information unavailable."""
-
-    instruction_text = "Medical Context:\n"
-    query_footer = f"\nQ: {rewritten_query}\nA:"
-
-    # Calculate static overhead tokens
-    inst_tokens = len(tokenizer.encode(instruction_text, add_special_tokens=False))
-    query_tokens = len(tokenizer.encode(query_footer, add_special_tokens=False))
-    total_static_cost = inst_tokens + query_tokens + 5
-
-    remaining_tokens = max(0, max_tokens - total_static_cost)
-
-    # Fill context budget
-    valid_contexts = []
-    current_context_tokens = 0
-
-    for idx, c in enumerate(context, start=1):
-        chunk_text = f"[C{idx}] {c['question']}\n{c['chunk_answer']}"
-        chunk_len = len(tokenizer.encode(chunk_text, add_special_tokens=False))
-
-        if current_context_tokens + chunk_len > remaining_tokens:
-            break
-
-        valid_contexts.append(chunk_text)
-        current_context_tokens += chunk_len
-
-    the_context_block = "\n".join(valid_contexts)
-    full_prompt = f"{instruction_text}{the_context_block}{query_footer}"
-
-    return full_prompt
-
-
-# ===========================
-# RESPONSE REFINEMENT
-# ===========================
-
-def refine_response(answer):
-    """Clean and format generated response text.
-    
-    Args:
-        answer (str): Raw generated text
-        
-    Returns:
-        str: Cleaned and formatted response
-    """
-    # Replace multiple periods
-    answer = re.sub(r'\. {2,}', '.', answer)
-    answer = re.sub(r'\.([^\s])', r'. \1', answer)
-
-    # Handle truncation at last punctuation
-    if not answer.strip().endswith(('.', '!', '?')):
-        last_punc_pos = max(answer.rfind('.'), answer.rfind('!'), answer.rfind('?'))
-        if last_punc_pos != -1:
-            answer = answer[:last_punc_pos + 1]
-
-    # Capitalize sentences
-    sentences = re.split(r'([.!?]\s*)', answer)
-    refined_sentences = []
-    for i in range(0, len(sentences), 2):
-        sentence_part = sentences[i].strip()
-        if sentence_part:
-            refined_sentences.append(sentence_part.capitalize())
-        if i + 1 < len(sentences):
-            refined_sentences.append(sentences[i + 1])
-
-    refined_sentences = ''.join(refined_sentences).strip()
-    return refined_sentences
-
-
-# ===========================
-# RAG PIPELINE
-# ===========================
-
-def rag_pipeline(user_query, top_k=3, detail=False):
-    """Main RAG pipeline: retrieve context and generate answer.
-    
-    Args:
-        user_query (str): User's medical question
-        top_k (int): Number of context chunks to retrieve
-        detail (bool): Whether to show detailed context information
-        
-    Returns:
-        str or dict: Answer with optional context details
-    """
-    try:
-        # Import retriever here to avoid circular imports
-        from retriever import retriever_simple
-        
-        # Check for emergencies
-        emergency_keywords = [
-            "emergency", "severe pain", "bleeding", "blind",
-            "lose consciousness", "pass out", "call 911"
-        ]
-        if any(keyword in user_query.lower() for keyword in emergency_keywords):
-            return """I am an AI and cannot provide medical advice for emergencies.
-PLEASE CONTACT EMERGENCY SERVICES OR A MEDICAL PROFESSIONAL IMMEDIATELY."""
-
-        # 1. Retrieve context
-        print(f"[RAG] Retrieving context for query: {user_query[:50]}...")
-        contexts = retriever_simple(user_query, top_k=top_k, detail=detail)
-        
-        if not contexts:
-            return "I couldn't find relevant medical information to answer your question. Please try rephrasing."
-
-        # 2. Build prompt
-        print(f"[RAG] Building prompt...")
-        prompt = build_prompt(user_query, contexts, user_query)
-
-        # 3. Generate response
-        print(f"[RAG] Generating response...")
-        llm = load_finetuned_model()
-        result = llm(
-            prompt,
-            max_new_tokens=70,
-            num_beams=3,
-            early_stopping=True,
-            do_sample=False,
-            repetition_penalty=1.4,
-            eos_token_id=llm.tokenizer.eos_token_id
-        )
-        
-        answer = result[0]['generated_text'].strip()
-        answer = refine_response(answer)
-
-        # 4. Calculate confidence from retrieval distances
-        if contexts and len(contexts) > 0:
-            avg_distance = sum(
-                c.get('chroma_distance', 1.0) for c in contexts
-            ) / len(contexts)
-            confidence_score = max(0, min(100, (1 - avg_distance) * 100))
-
-            if confidence_score < 40:
-                final_answer = f"⚠️ Low Confidence ({confidence_score:.1f}%)\n\n{answer}"
-            else:
-                final_answer = f"{answer}\n\n✓ Confidence: {confidence_score:.1f}%"
-        else:
-            final_answer = answer
-
-        return final_answer
-
-    except Exception as e:
-        error_msg = f"Error in RAG pipeline: {str(e)}"
-        print(error_msg)
-        traceback.print_exc()
-        return error_msg