Delete retriever.py

retriever.py

@@ -1,354 +0,0 @@
-"""
-Retrieval functions for the RAG pipeline.
-Handles query rewriting, retrieval, and context re-ranking.
-"""
-
-import torch
-from model import (
-    load_embeddings_model,
-    load_vector_store,
-    load_rewriter_model,
-    load_reranker_model
-)
-
-
-# ===========================
-# QUERY REWRITING
-# ===========================
-
-def rewrite_query(user_query):
-    """Rewrite user query to be more specific and medical-focused.
-    
-    Args:
-        user_query (str): Original user question
-        
-    Returns:
-        str: Rewritten query
-    """
-    rewriter_llm = load_rewriter_model()
-
-    # Few-shot prompting for medical question reformulation
-    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 = rewriter_llm(prompt)
-    rewritten_query = llm_output[0]['generated_text']
-    rewritten_query = rewritten_query.replace("Output:", "").strip()
-
-    return rewritten_query.strip()
-
-
-# ===========================
-# RERANKING
-# ===========================
-
-def get_monot5_scores(rewritten_query, chunks):
-    """Calculate relevance scores for chunks using MonoT5 reranker.
-    
-    Args:
-        rewritten_query (str): The rewritten query
-        chunks (list): List of retrieved text chunks
-        
-    Returns:
-        list: Relevance scores for each chunk
-    """
-    rerank_tokenizer, rerank_model = load_reranker_model()
-
-    # Format inputs: "Query: Q Document: D 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 scores
-        true_token_id = rerank_tokenizer.encode("true")[0]
-        batch_scores = outputs.scores[0][:, true_token_id]
-
-        return batch_scores.tolist()
-
-
-# ===========================
-# RETRIEVAL FUNCTIONS
-# ===========================
-
-def retriever_simple(q, top_k=3, detail=False):
-    """Simple retrieval without query rewriting or reranking.
-    
-    Args:
-        q (str or list): Query or list of queries
-        top_k (int): Number of results to return
-        detail (bool): Include metadata in results
-        
-    Returns:
-        list: Retrieved context chunks
-    """
-    embed_model = load_embeddings_model()
-    vector_store = load_vector_store()
-
-    is_single_query = isinstance(q, str)
-    queries = [q] if is_single_query else q
-
-    # Encode queries
-    q_embeddings = embed_model.encode(queries).tolist()
-
-    # Search vector store
-    search_results = vector_store.query(
-        query_embeddings=q_embeddings,
-        n_results=10
-    )
-
-    all_contexts = []
-
-    for i in range(len(queries)):
-        contexts_for_query = []
-        if search_results['documents'][i]:
-            for r in range(len(search_results['ids'][i])):
-                item = {
-                    'chunk_id': search_results['ids'][i][r],
-                    'chunk_answer': search_results['documents'][i][r],
-                    'question': search_results['metadatas'][i][r]['question']
-                }
-                if detail:
-                    item['focus_area'] = search_results['metadatas'][i][r].get('focus_area', 'Unknown')
-                    item['source'] = search_results['metadatas'][i][r].get('source', 'Unknown')
-                    item['chroma_distance'] = round(search_results['distances'][i][r], 3)
-                contexts_for_query.append(item)
-        all_contexts.append(contexts_for_query)
-
-    if is_single_query:
-        return all_contexts[0]
-    else:
-        return all_contexts
-
-
-def retriever_with_rewriter(q, top_k=3, detail=False):
-    """Retrieve with query rewriting but without reranking.
-    
-    Args:
-        q (str or list): Query or list of queries
-        top_k (int): Number of results to return
-        detail (bool): Include metadata in results
-        
-    Returns:
-        list: Retrieved context chunks
-    """
-    embed_model = load_embeddings_model()
-    vector_store = load_vector_store()
-
-    is_single_query = isinstance(q, str)
-    queries = [q] if is_single_query else q
-
-    # Rewrite queries
-    rewritten_queries = [rewrite_query(query) for query in queries]
-
-    # Encode rewritten queries
-    q_embeddings = embed_model.encode(rewritten_queries).tolist()
-
-    # Search vector store
-    search_results = vector_store.query(
-        query_embeddings=q_embeddings,
-        n_results=10
-    )
-
-    all_contexts = []
-
-    for i in range(len(queries)):
-        contexts_for_query = []
-        if search_results['documents'][i]:
-            for r in range(min(top_k, len(search_results['ids'][i]))):
-                item = {
-                    'chunk_id': search_results['ids'][i][r],
-                    'chunk_answer': search_results['documents'][i][r],
-                    'question': search_results['metadatas'][i][r]['question']
-                }
-                if detail:
-                    item['focus_area'] = search_results['metadatas'][i][r].get('focus_area', 'Unknown')
-                    item['source'] = search_results['metadatas'][i][r].get('source', 'Unknown')
-                    item['chroma_distance'] = round(search_results['distances'][i][r], 3)
-                contexts_for_query.append(item)
-        all_contexts.append(contexts_for_query)
-
-    if is_single_query:
-        return all_contexts[0]
-    else:
-        return all_contexts
-
-
-def retriever_with_reranker(q, top_k=3, detail=False):
-    """Retrieve with reranking but without query rewriting.
-    
-    Args:
-        q (str or list): Query or list of queries
-        top_k (int): Number of results to return
-        detail (bool): Include metadata in results
-        
-    Returns:
-        list: Retrieved context chunks
-    """
-    embed_model = load_embeddings_model()
-    vector_store = load_vector_store()
-
-    is_single_query = isinstance(q, str)
-    queries = [q] if is_single_query else q
-
-    q_embeddings = embed_model.encode(queries).tolist()
-
-    search_results = vector_store.query(
-        query_embeddings=q_embeddings,
-        n_results=10
-    )
-
-    all_contexts = []
-
-    for i in range(len(queries)):
-        contexts_for_query = []
-        if search_results['documents'][i]:
-            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
-            rerank_scores = get_monot5_scores(queries[i], retrieved_chunks)
-
-            if len(rerank_scores) == len(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_distance': retrieved_distances[r],
-                        'rerank_score': rerank_scores[r]
-                    })
-
-                # Sort by rerank score
-                candidates.sort(key=lambda x: x['rerank_score'], reverse=True)
-                sorted_results = candidates[:top_k]
-
-                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['rerank_score'] = round(item_dict['rerank_score'], 3)
-                        item['chroma_distance'] = round(item_dict['chroma_distance'], 3)
-                    contexts_for_query.append(item)
-
-        all_contexts.append(contexts_for_query)
-
-    if is_single_query:
-        return all_contexts[0]
-    else:
-        return all_contexts
-
-
-def retriever_full(q, top_k=3, detail=False):
-    """Full retrieval with both query rewriting and reranking (recommended).
-    
-    Args:
-        q (str or list): Query or list of queries
-        top_k (int): Number of results to return
-        detail (bool): Include metadata in results
-        
-    Returns:
-        list: Retrieved context chunks
-    """
-    embed_model = load_embeddings_model()
-    vector_store = load_vector_store()
-
-    is_single_query = isinstance(q, str)
-    queries = [q] if is_single_query else q
-
-    all_contexts = []
-
-    for query in queries:
-        # Rewrite query
-        q_rewritten = rewrite_query(query)
-        q_embed = embed_model.encode([q_rewritten]).tolist()
-
-        # Search
-        search_result = vector_store.query(
-            query_embeddings=q_embed,
-            n_results=10
-        )
-
-        if not search_result['documents'][0]:
-            all_contexts.append([])
-            continue
-
-        retrieved_chunks = search_result['documents'][0]
-        retrieved_metas = search_result['metadatas'][0]
-        retrieved_ids = search_result['ids'][0]
-        retrieved_distances = search_result.get('distances', [[]])[0]
-
-        # Rerank
-        rerank_scores = get_monot5_scores(q_rewritten, retrieved_chunks)
-
-        if len(rerank_scores) == len(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_distance': retrieved_distances[r],
-                    'rerank_score': rerank_scores[r]
-                })
-
-            # Sort by rerank score
-            candidates.sort(key=lambda x: x['rerank_score'], reverse=True)
-            sorted_results = candidates[:top_k]
-
-            contexts = []
-            for i in sorted_results:
-                item = {
-                    'chunk_id': i['chunk_id'],
-                    'chunk_answer': i['chunk_answer'],
-                    'question': i['question']
-                }
-                if detail:
-                    item['focus_area'] = i['focus_area']
-                    item['rerank_score'] = round(i['rerank_score'], 3)
-                    item['chroma_distance'] = round(i['chroma_distance'], 3)
-                contexts.append(item)
-
-            all_contexts.append(contexts)
-        else:
-            all_contexts.append([])
-
-    if is_single_query:
-        return all_contexts[0]
-    else:
-        return all_contexts