Update retriever.py

retriever.py

@@ -1,185 +1,185 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from utils import get_embed_model
-
-
-# ============================================================================
-# LAZY-LOAD RERANKER AND QUERY REWRITER (load on first use)
-# ============================================================================
-
-_rerank_tokenizer = None
-_rerank_model = None
-_rewritter_llm = None
-
-def get_reranker():
-    """Load reranker model (lazy-loaded on first use)"""
-    global _rerank_tokenizer, _rerank_model
-    if _rerank_tokenizer is None:
-        _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 get_query_rewriter():
-    """Load query rewriter (lazy-loaded on first use)"""
-    global _rewritter_llm
-    if _rewritter_llm is None:
-        _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
-        )
-    return _rewritter_llm
-
-
-def rewrite_query(user_query):
-    """Rewrite user query to be more specific and medical-focused.
-    Returns: str(rewritten_query)"""
-    
-    rewritter_llm = get_query_rewriter()
-    
-    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)"""
-    
-    rerank_tokenizer, rerank_model = get_reranker()
-    
-    # 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)"""
-    
-    embed_model = get_embed_model()
-    
-    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
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from utils import get_embed_model
+
+
+# ============================================================================
+# LAZY-LOAD RERANKER AND QUERY REWRITER (load on first use)
+# ============================================================================
+
+_rerank_tokenizer = None
+_rerank_model = None
+_rewritter_llm = None
+
+def get_reranker():
+    """Load reranker model (lazy-loaded on first use)"""
+    global _rerank_tokenizer, _rerank_model
+    if _rerank_tokenizer is None:
+        _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 get_query_rewriter():
+    """Load query rewriter (lazy-loaded on first use)"""
+    global _rewritter_llm
+    if _rewritter_llm is None:
+        _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
+        )
+    return _rewritter_llm
+
+
+def rewrite_query(user_query):
+    """Rewrite user query to be more specific and medical-focused.
+    Returns: str(rewritten_query)"""
+    
+    rewritter_llm = get_query_rewriter()
+    
+    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)"""
+    
+    rerank_tokenizer, rerank_model = get_reranker()
+    
+    # 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)"""
+    
+    embed_model = get_embed_model()
+    
+    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