Upload 5 files

app.py

@@ -0,0 +1,13 @@
+import gradio as gr
+from model import rag
+
+demo = gr.Interface(
+    fn = rag,
+    inputs = gr.Textbox(lines=5, label="Enter your medical question here"),
+    outputs = gr.Textbox(lines=10, label="Generated Answer"),
+    title = "Medical QA",
+    description = "Type your query and get a model generated response.",
+    theme = gr.themes.Glass()
+)
+
+demo.launch()

model.py

@@ -0,0 +1,130 @@
+import chromadb
+import traceback
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from retriever import retrieve
+from utils import build_prompt, refine_response
+
+
+# ============================================================================
+# LOAD MODELS AND VECTOR STORE ONCE AT STARTUP
+# ============================================================================
+
+# Load vector store
+db_client = chromadb.PersistentClient(path="./MedQuAD_db")
+try:
+    vector_store = db_client.get_collection("medical_rag")
+except:
+    # If collection doesn't exist, create it
+    vector_store = db_client.create_collection(name="medical_rag")
+
+# Load fine-tuned model for generation
+ft_model_id = "amiraghhh/fine-tuned-flan-t5-small"
+ft_tokenizer = AutoTokenizer.from_pretrained(ft_model_id)
+ft_model = AutoModelForSeq2SeqLM.from_pretrained(ft_model_id)
+
+finetuned_llm = pipeline(
+    "text2text-generation",
+    model=ft_model,
+    tokenizer=ft_tokenizer,
+    decoder_start_token_id=ft_model.config.pad_token_id
+)
+
+# Load base model config for EOS token
+base_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+
+
+# ============================================================================
+# MAIN RAG FUNCTION
+# ============================================================================
+
+def rag(user_query):
+    """Main RAG function: retrieve context and generate answer.
+    Takes a question string and returns an answer string with confidence.
+    Returns: str(generated_answer)"""
+    
+    try:
+        # 1. Check for emergency keywords
+        emergency_keywords = ["emergency", "severe pain", "bleeding",
+                            "blind", "lose consciousness", "pass out"]
+        
+        if any(keyword in user_query.lower() for keyword in emergency_keywords):
+            emergency_msg = """I am an AI and cannot provide medical advice for emergencies.
+PLEASE contact emergency services or a medical professional immediately."""
+            
+            try:
+                # Still generate answer for context
+                contexts = retrieve(vector_store, user_query, top_k=3, use_reranking=True)
+                
+                if not contexts:
+                    return f"{emergency_msg}\n\nNo relevant information found for your query."
+                
+                prompt = build_prompt(user_query, contexts)
+                result = finetuned_llm(
+                    prompt,
+                    max_new_tokens=70,
+                    num_beams=3,
+                    early_stopping=True,
+                    do_sample=False,
+                    repetition_penalty=1.4,
+                    eos_token_id=finetuned_llm.tokenizer.eos_token_id
+                )
+                
+                answer = result[0]['generated_text'].strip()
+                answer = refine_response(answer)
+                
+                # Calculate confidence
+                if contexts:
+                    avg_distance = sum(c.get('chroma_distance', 1.0) for c in contexts) / len(contexts)
+                    confidence_score = (1 - avg_distance) * 100
+                    confidence_score = max(0, min(100, confidence_score))
+                else:
+                    confidence_score = 0
+                
+                return f"{emergency_msg}\n\n[Confidence: {confidence_score:.1f}%]\n\n{answer}"
+            
+            except Exception as e:
+                return f"{emergency_msg}\n\nError generating answer: {str(e)}"
+        
+        # 2. Retrieve relevant contexts
+        contexts = retrieve(vector_store, user_query, top_k=3, use_reranking=True)
+        
+        if not contexts:
+            return "I'm not confident about my answer (0%).\n\nCouldn't find relevant information to answer your question."
+        
+        # 3. Build prompt with context
+        prompt = build_prompt(user_query, contexts)
+        
+        # 4. Generate answer
+        result = finetuned_llm(
+            prompt,
+            max_new_tokens=70,
+            num_beams=3,
+            early_stopping=True,
+            do_sample=False,
+            repetition_penalty=1.4,
+            eos_token_id=finetuned_llm.tokenizer.eos_token_id
+        )
+        
+        answer = result[0]['generated_text'].strip()
+        answer = refine_response(answer)
+        
+        # 5. Calculate confidence score based on retrieval quality
+        if contexts and len(contexts) > 0:
+            avg_distance = sum(c.get('chroma_distance', 1.0) for c in contexts) / len(contexts)
+            confidence_score = (1 - avg_distance) * 100
+            confidence_score = max(0, min(100, confidence_score))
+            
+            # Build final response with confidence
+            if confidence_score < 40:
+                final_response = f"I'm not confident about my answer ({confidence_score:.1f}%).\n\n{answer}"
+            else:
+                final_response = f"{answer}\n\n[Confidence: {confidence_score:.1f}%]"
+        else:
+            final_response = "I'm not confident about my answer (0%).\n\n" + answer
+        
+        return final_response
+    
+    except Exception as e:
+        error_msg = f"ERROR in RAG pipeline: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+        print(error_msg)
+        return error_msg

requirements.txt

@@ -0,0 +1,9 @@
+gradio==4.26.0
+torch==2.1.2
+transformers==4.37.2
+sentence-transformers==2.2.2
+chromadb==0.4.24
+langchain-text-splitters==0.0.1
+accelerate==0.27.2
+numpy==1.24.3
+pandas==2.0.3

retriever.py

@@ -0,0 +1,163 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from utils import embed_model
+
+
+# Initialize reranker model for relevance scoring
+rerank_tokenizer = AutoTokenizer.from_pretrained("castorini/monot5-base-msmarco")
+rerank_model = AutoModelForSeq2SeqLM.from_pretrained("castorini/monot5-base-msmarco")
+rerank_model.eval()
+
+# Initialize query rewriter
+rewritter_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
+)
+
+
+def rewrite_query(user_query):
+    """Rewrite user query to be more specific and medical-focused.
+    Returns: str(rewritten_query)"""
+    
+    prompt = f"""Rewrite the input into a clear medical question following these patterns
+    Input: my head hurts
+    Output: What causes headaches?
+
+    Input: i keep vomiting but feel ok afterwards
+    Output: What causes cyclic vomiting?
+
+    Input: chest pain when breathing
+    Output: What causes chest pain during breathing?
+
+    Input: {user_query}
+    Output:
+    """
+    
+    llm_output = rewritter_llm(prompt)
+    rewritten_query = llm_output[0]['generated_text']
+    rewritten_query = rewritten_query.replace("Output:", "").strip()
+    
+    return rewritten_query.strip()
+
+
+def get_monot5_scores(rewritten_query, chunks):
+    """Calculate relevance scores for chunks against query using MonoT5 reranker.
+    Returns: list(scores)"""
+    
+    # Format input as "Query: ABC? Document: XYZ... Relevant:"
+    inputs = [f"Query: {rewritten_query} Document: {c} Relevant:" for c in chunks]
+    
+    # Tokenize
+    tokenized_inputs = rerank_tokenizer(
+        inputs,
+        padding=True,
+        truncation=True,
+        return_tensors='pt'
+    )
+    
+    # Generate predictions
+    with torch.no_grad():
+        outputs = rerank_model.generate(
+            input_ids=tokenized_inputs['input_ids'],
+            attention_mask=tokenized_inputs['attention_mask'],
+            max_new_tokens=1,
+            return_dict_in_generate=True,
+            output_scores=True
+        )
+        
+        # Extract "true" token probability scores
+        true_token_id = rerank_tokenizer.encode("true")[0]
+        batch_scores = outputs.scores[0][:, true_token_id]
+        
+        return batch_scores.tolist()
+
+
+def retrieve(vector_store, query, top_k=3, use_reranking=True, detail=False):
+    """Retrieve relevant context chunks for a query with optional reranking.
+    Returns: list(contexts)"""
+    
+    is_single_query = isinstance(query, str)
+    queries = [query] if is_single_query else query
+    
+    # Rewrite queries for better matching
+    rewritten_queries = [rewrite_query(q) for q in queries]
+    
+    # Embed rewritten queries
+    q_embeddings = embed_model.encode(rewritten_queries).tolist()
+    
+    # Vector search
+    search_results = vector_store.query(
+        query_embeddings=q_embeddings,
+        n_results=10
+    )
+    
+    all_contexts = []
+    
+    for i in range(len(queries)):
+        contexts_for_query = []
+        
+        if not search_results['documents'][i]:
+            all_contexts.append([])
+            continue
+        
+        retrieved_chunks = search_results['documents'][i]
+        retrieved_metas = search_results['metadatas'][i]
+        retrieved_ids = search_results['ids'][i]
+        retrieved_distances = search_results.get('distances', [[]])[i]
+        
+        # Rerank if enabled
+        if use_reranking:
+            rerank_scores = get_monot5_scores(rewritten_queries[i], retrieved_chunks)
+            
+            candidates = []
+            for r in range(len(retrieved_chunks)):
+                candidates.append({
+                    'chunk_id': retrieved_ids[r],
+                    'chunk_answer': retrieved_chunks[r],
+                    'question': retrieved_metas[r]['question'],
+                    'focus_area': retrieved_metas[r].get('focus_area', 'Unknown'),
+                    'chroma_dist': retrieved_distances[r],
+                    'rerank_score': rerank_scores[r]
+                })
+            
+            # Sort by rerank score
+            candidates.sort(key=lambda s: s['rerank_score'], reverse=True)
+            sorted_results = candidates[:top_k]
+        else:
+            # Just use top-k from vector search
+            sorted_results = [
+                {
+                    'chunk_id': retrieved_ids[r],
+                    'chunk_answer': retrieved_chunks[r],
+                    'question': retrieved_metas[r]['question'],
+                    'focus_area': retrieved_metas[r].get('focus_area', 'Unknown'),
+                    'chroma_dist': retrieved_distances[r]
+                }
+                for r in range(min(top_k, len(retrieved_chunks)))
+            ]
+        
+        # Format output
+        for item_dict in sorted_results:
+            item = {
+                'chunk_id': item_dict['chunk_id'],
+                'chunk_answer': item_dict['chunk_answer'],
+                'question': item_dict['question']
+            }
+            if detail:
+                item['focus_area'] = item_dict['focus_area']
+                item['chroma_distance'] = round(item_dict['chroma_dist'], 3)
+                if 'rerank_score' in item_dict:
+                    item['rerank_score'] = round(item_dict['rerank_score'], 3)
+            contexts_for_query.append(item)
+        
+        all_contexts.append(contexts_for_query)
+    
+    if is_single_query:
+        return all_contexts[0]
+    else:
+        return all_contexts

utils.py

@@ -0,0 +1,114 @@
+import re
+from hashlib import md5
+from sentence_transformers import SentenceTransformer
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from transformers import AutoTokenizer
+
+# Initialize embedding model (used across multiple modules)
+embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+flant5tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+
+
+def normalize_text(text):
+    """Normalize text for duplicate detection by removing spaces and punctuation.
+    Returns: str"""
+    if not isinstance(text, str):
+        return ""
+    
+    text = text.lower()  # Lowercase
+    text = re.sub(r'\s+', ' ', text)  # Remove extra whitespace
+    text = text.strip()  # Remove leading/trailing spaces
+    text = re.sub(r'[^\w\s]', '', text)  # Remove punctuation
+    
+    return text
+
+
+def chunk_text(text, chunk_size=384, chunk_overlap=20):
+    """Split text into chunks for embeddings.
+    Returns: list(chunks)"""
+    splitter = RecursiveCharacterTextSplitter(
+        separators=["\n\n", "\n", " ", ""],
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap
+    )
+    return splitter.split_text(text)
+
+
+def create_embeddings(texts):
+    """Create embeddings for a list of texts.
+    Returns: list(embeddings)"""
+    return embed_model.encode(
+        texts,
+        batch_size=64,
+        show_progress_bar=False,
+        convert_to_numpy=True,
+        normalize_embeddings=True
+    )
+
+
+def refine_response(answer):
+    """Clean and format generated response text.
+    Returns: str(refined_answer)"""
+    # Replace multiple periods with single period
+    answer = re.sub(r'\. {2,}', '.', answer)
+    answer = re.sub(r'\.([^\s])', r'. \1', answer)
+    
+    # If no ending punctuation, find last one and truncate
+    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 first letter of each sentence
+    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])
+    
+    return ''.join(refined_sentences).strip()
+
+
+def build_prompt(user_query, context, max_tokens=512):
+    """Build prompt with context and query within token limit.
+    Returns: str(full_prompt)"""
+    
+    if not context:
+        return f"""No relevant medical information found.
+Q: {user_query}
+A: Information unavailable."""
+    
+    instruction_text = "Medical Context:\n"
+    query_footer = f"\nQ: {user_query}\nA:"
+    
+    # Calculate token overhead
+    inst_tokens = len(flant5tokenizer.encode(instruction_text, add_special_tokens=False))
+    query_tokens = len(flant5tokenizer.encode(query_footer, add_special_tokens=False))
+    total_static_cost = inst_tokens + query_tokens + 5
+    
+    # Calculate remaining budget for context
+    remaining_tokens = max_tokens - total_static_cost
+    if remaining_tokens < 0:
+        remaining_tokens = 0
+    
+    # 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(flant5tokenizer.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