Advanced Retrieval-Augmented Generation with LLaMA 3: Beyond Basic RAG

def decompose_query(query):
    decomposition_prompt = f"""<|system|>
You are an expert at breaking down complex questions into simpler sub-questions.
</s>
<|user|>
Break down the following question into 2-4 sequential sub-questions that would help answer the original question when combined:

Question: {query}
</s>
<|assistant|>
"""
    # Generate decomposition using LLaMA 3
    response = llm.generate(decomposition_prompt)
    
    # Parse sub-questions from response
    sub_questions = [q.strip() for q in response.split("\n") if q.strip().startswith("Sub-question")]
    return sub_questions

def hybrid_retrieval(query, top_k=5):
    # Generate dense embeddings for semantic search
    query_embedding = embedding_model.embed_query(query)
    dense_results = vector_store.similarity_search_by_vector(query_embedding, k=top_k*2)
    
    # Extract key entities and terms for sparse search
    key_terms = extract_key_terms(query, llm)
    sparse_results = keyword_index.search(key_terms, k=top_k*2)
    
    # Combine results with a re-ranking step
    combined_results = list(set(dense_results + sparse_results))
    reranked_results = rerank_results(query, combined_results, llm)
    
    return reranked_results[:top_k]

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import List, Dict, Any
import numpy as np

# Initialize LLaMA 3 model
def initialize_llm(model_path="meta-llama/Llama-3-8B-Instruct"):
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None
    )
    return model, tokenizer

# Semantic Document Chunker
class SemanticChunker:
    def __init__(self, llm, tokenizer):
        self.llm = llm
        self.tokenizer = tokenizer
        
    def identify_section_boundaries(self, document: str) -> List[Dict[str, Any]]:
        """Identify logical section boundaries in a document"""
        boundary_prompt = f"""<|system|>
Analyze the following document and identify its logical sections. 
For each section, provide:
1. The section title or topic
2. The starting line number
3. A brief description of what the section covers
</s>
<|user|>
{document[:5000]}  # First portion of the document
</s>
<|assistant|>
"""
        inputs = self.tokenizer(boundary_prompt, return_tensors="pt").to(self.llm.device)
        with torch.no_grad():
            outputs = self.llm.generate(
                **inputs,
                max_new_tokens=1024,
                temperature=0.1,
            )
        
        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        
        # Parse the sections from the response (implementation omitted for brevity)
        sections = self._parse_sections(response, document)
        return sections
    
    def chunk_document(self, document: str, max_chunk_size: int = 512) -> List[Dict[str, Any]]:
        """Chunk document respecting semantic boundaries"""
        sections = self.identify_section_boundaries(document)
        chunks = []
        
        for section in sections:
            section_text = document[section["start_idx"]:section["end_idx"]]
            
            # If section fits within chunk size, keep it whole
            if len(self.tokenizer.encode(section_text)) <= max_chunk_size:
                chunks.append({
                    "text": section_text,
                    "metadata": {
                        "section": section["title"],
                        "summary": section["description"]
                    }
                })
            else:
                # For larger sections, split while trying to respect paragraph boundaries
                sub_chunks = self._split_section(section_text, max_chunk_size)
                for i, sub_chunk in enumerate(sub_chunks):
                    chunks.append({
                        "text": sub_chunk,
                        "metadata": {
                            "section": section["title"],
                            "part": i+1,
                            "summary": section["description"]
                        }
                    })
        
        return chunks
    
    def _parse_sections(self, response, document):
        # Implementation omitted for brevity
        # This would parse the model's output to extract section information
        pass
    
    def _split_section(self, section_text, max_chunk_size):
        # Implementation omitted for brevity
        # This would split a section into smaller chunks respecting paragraph boundaries
        pass

from sentence_transformers import SentenceTransformer, CrossEncoder

class EnhancedEmbeddingGenerator:
    def __init__(self, bi_encoder_model="BAAI/bge-large-en-v1.5", cross_encoder_model="cross-encoder/ms-marco-MiniLM-L-6-v2"):
        self.bi_encoder = SentenceTransformer(bi_encoder_model)
        self.cross_encoder = CrossEncoder(cross_encoder_model)
        
    def generate_embeddings(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Generate embeddings for document chunks"""
        texts = [chunk["text"] for chunk in chunks]
        embeddings = self.bi_encoder.encode(texts, normalize_embeddings=True)
        
        for i, chunk in enumerate(chunks):
            chunk["embedding"] = embeddings[i]
            
            # Generate summary embedding focusing on key concepts
            if "summary" in chunk["metadata"]:
                summary_embedding = self.bi_encoder.encode(chunk["metadata"]["summary"], normalize_embeddings=True)
                # Store both text and summary embeddings
                chunk["summary_embedding"] = summary_embedding
        
        return chunks
    
    def score_relevance(self, query: str, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Score relevance of chunks to query using cross-encoder"""
        texts = [chunk["text"] for chunk in chunks]
        pairs = [[query, text] for text in texts]
        
        # Calculate relevance scores
        relevance_scores = self.cross_encoder.predict(pairs)
        
        # Add scores to chunks
        for i, chunk in enumerate(chunks):
            chunk["relevance_score"] = float(relevance_scores[i])
            
        return chunks

class AdvancedRAGSystem:
    def __init__(self, llm, tokenizer, embedding_generator, vector_store):
        self.llm = llm
        self.tokenizer = tokenizer
        self.embedding_generator = embedding_generator
        self.vector_store = vector_store
        
    def answer_query(self, query: str) -> Dict[str, Any]:
        # Step 1: Decompose complex queries
        sub_queries = self._decompose_query(query) if self._is_complex_query(query) else [query]
        
        all_retrieved_chunks = []
        # Step 2: Retrieve relevant chunks for each sub-query
        for sub_query in sub_queries:
            retrieved_chunks = self._retrieve_chunks(sub_query)
            all_retrieved_chunks.extend(retrieved_chunks)
        
        # Step 3: Remove duplicates and rerank all retrieved chunks
        unique_chunks = self._deduplicate_chunks(all_retrieved_chunks)
        reranked_chunks = self.embedding_generator.score_relevance(query, unique_chunks)
        top_chunks = sorted(reranked_chunks, key=lambda x: x["relevance_score"], reverse=True)[:5]
        
        # Step 4: Generate initial response
        initial_response = self._generate_response(query, top_chunks)
        
        # Step 5: Self-critique and refine
        final_response = self._refine_response(query, initial_response, top_chunks)
        
        return {
            "query": query,
            "response": final_response,
            "supporting_chunks": top_chunks,
            "sub_queries": sub_queries
        }
    
    def _is_complex_query(self, query: str) -> bool:
        # Determine if query needs decomposition
        complexity_prompt = f"""<|system|>
Determine if the following query is complex and would benefit from being broken down into sub-questions.
A complex query typically involves multiple concepts, comparisons, or logical steps.
</s>
<|user|>
Query: {query}
Is this a complex query that should be decomposed? Answer with Yes or No and a brief explanation.
</s>
<|assistant|>
"""
        inputs = self.tokenizer(complexity_prompt, return_tensors="pt").to(self.llm.device)
        with torch.no_grad():
            outputs = self.llm.generate(
                **inputs,
                max_new_tokens=100,
                temperature=0.1,
            )
        
        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        return "yes" in response.lower()
    
    def _decompose_query(self, query: str) -> List[str]:
        # Implementation details omitted for brevity
        # This would use LLaMA 3 to break down the query into sub-queries
        pass
    
    def _retrieve_chunks(self, query: str) -> List[Dict[str, Any]]:
        # Vector similarity search plus optional keyword filtering
        query_embedding = self.embedding_generator.bi_encoder.encode(query, normalize_embeddings=True)
        retrieved_chunks = self.vector_store.similarity_search_by_vector(query_embedding)
        return retrieved_chunks
    
    def _deduplicate_chunks(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        # Remove duplicate chunks based on content similarity
        # Implementation details omitted for brevity
        pass
    
    def _generate_response(self, query: str, chunks: List[Dict[str, Any]]) -> str:
        # Construct context from chunks
        context = "\n\n".join([f"Document {i+1}:\n{chunk['text']}" for i, chunk in enumerate(chunks)])
        
        response_prompt = f"""<|system|>
You are a knowledgeable assistant that responds to questions based on the provided context.
When answering, cite the specific document numbers that contain supporting information.
If the context doesn't contain relevant information, acknowledge the limitations of your knowledge.
</s>
<|user|>
Context:
{context}

Question: {query}
</s>
<|assistant|>
"""
        
        inputs = self.tokenizer(response_prompt, return_tensors="pt").to(self.llm.device)
        with torch.no_grad():
            outputs = self.llm.generate(
                **inputs,
                max_new_tokens=512,
                temperature=0.7,
                top_p=0.9,
            )
        
        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        return response
    
    def _refine_response(self, query: str, initial_response: str, chunks: List[Dict[str, Any]]) -> str:
        # Critique and refine the initial response
        critique_prompt = f"""<|system|>
You are a critical evaluator assessing the quality of an AI-generated response.
Check for:
1. Factual accuracy compared to the source documents
2. Completeness in addressing all aspects of the query
3. Proper citations to supporting evidence
4. Clarity and coherence
</s>
<|user|>
Query: {query}

Response to evaluate:
{initial_response}

Source documents:
{"\n\n".join([f"Document {i+1}:\n{chunk['text']}" for i, chunk in enumerate(chunks)])}

Provide a critique highlighting any issues found.
</s>
<|assistant|>
"""
        
        inputs = self.tokenizer(critique_prompt, return_tensors="pt").to(self.llm.device)
        with torch.no_grad():
            outputs = self.llm.generate(
                **inputs,
                max_new_tokens=256,
                temperature=0.3,
            )
        
        critique = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        
        # Generate refined response based on critique
        refine_prompt = f"""<|system|>
You are a knowledgeable assistant that creates high-quality responses based on provided source documents.
</s>
<|user|>
Query: {query}

Your previous response:
{initial_response}

Critique of your response:
{critique}

Source documents:
{"\n\n".join([f"Document {i+1}:\n{chunk['text']}" for i, chunk in enumerate(chunks)])}

Please provide an improved response that addresses the issues identified in the critique.
</s>
<|assistant|>
"""
        
        inputs = self.tokenizer(refine_prompt, return_tensors="pt").to(self.llm.device)
        with torch.no_grad():
            outputs = self.llm.generate(
                **inputs,
                max_new_tokens=512,
                temperature=0.7,
                top_p=0.9,
            )
        
        refined_response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        return refined_response

class RAGEvaluator:
    def __init__(self, llm, tokenizer, reference_data=None):
        self.llm = llm
        self.tokenizer = tokenizer
        self.reference_data = reference_data
        
    def evaluate_response(self, query, response, retrieved_chunks, ground_truth=None):
        results = {}
        
        # 1. Factual Accuracy
        results["factual_accuracy"] = self._evaluate_factual_accuracy(response, retrieved_chunks, ground_truth)
        
        # 2. Retrieval Precision
        results["retrieval_precision"] = self._evaluate_retrieval_precision(query, retrieved_chunks, ground_truth)
        
        # 3. Answer Completeness
        results["answer_completeness"] = self._evaluate_completeness(query, response, retrieved_chunks)
        
        # 4. Citation Quality
        results["citation_quality"] = self._evaluate_citations(response, retrieved_chunks)
        
        # 5. Hallucination Detection
        results["hallucination_score"] = self._detect_hallucinations(response, retrieved_chunks)
        
        # Calculate aggregate score
        results["aggregate_score"] = sum([
            results["factual_accuracy"] * 0.3,
            results["retrieval_precision"] * 0.2,
            results["answer_completeness"] * 0.2,
            results["citation_quality"] * 0.15,
            (1 - results["hallucination_score"]) * 0.15  # Lower hallucination is better
        ])
        
        return results
    
    def _evaluate_factual_accuracy(self, response, retrieved_chunks, ground_truth=None):
        # Implementation details omitted for brevity
        # Use LLaMA 3 to compare response claims against retrieved chunks
        pass
    
    def _evaluate_retrieval_precision(self, query, retrieved_chunks, ground_truth=None):
        # Implementation details omitted for brevity
        # Assess whether retrieved chunks contain information needed to answer query
        pass
    
    def _evaluate_completeness(self, query, response, retrieved_chunks):
        # Implementation details omitted for brevity
        # Determine if response addresses all aspects of the query
        pass
    
    def _evaluate_citations(self, response, retrieved_chunks):
        # Implementation details omitted for brevity
        # Check if claims in response are properly attributed to source documents
        pass
    
    def _detect_hallucinations(self, response, retrieved_chunks):
        # Implementation details omitted for brevity
        # Identify statements in response not supported by retrieved chunks
        pass