We Ditched RAG 2.0 for Fine-Tuned Llama 3.1: 50% Better Accuracy for Internal Docs

import os
import json
import logging
from typing import List, Dict, Any
from pathlib import Path
import unstructured.partition.auto as partition
from datasets import Dataset, DatasetDict
from sentence_transformers import SentenceTransformer
import numpy as np

# Configure logging for audit trails
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.FileHandler('data_prep.log'), logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# Constants for chunking and dataset formatting
MAX_CHUNK_SIZE = 1024
OVERLAP_SIZE = 128
OUTPUT_DATASET_PATH = Path('./fine_tuning_data')
LLAMA_3_1_INSTRUCTION_TEMPLATE = '''<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a technical documentation assistant for [Company Name Redacted] engineering team. Answer queries using only the provided context. If the answer is not in the context, say "I don't have enough information to answer that."<|eot_id|><|start_header_id|>user<|end_header_id|>
Context: {context}

Query: {query}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
{answer}'''

def validate_file_path(file_path: Path) -> None:
    '''Raise FileNotFoundError if input path is invalid.'''
    if not file_path.exists():
        raise FileNotFoundError(f'Input file {file_path} does not exist')
    if not file_path.is_file():
        raise ValueError(f'Input path {file_path} is not a file')

def chunk_document(text: str, max_chunk_size: int = MAX_CHUNK_SIZE, overlap: int = OVERLAP_SIZE) -> List[str]:
    '''Split long documents into overlapping chunks to preserve context.'''
    if len(text) <= max_chunk_size:
        return [text]
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + max_chunk_size, len(text))
        chunk = text[start:end]
        chunks.append(chunk)
        start += max_chunk_size - overlap
    return chunks

def process_internal_doc(doc_path: Path) -> List[Dict[str, Any]]:
    '''Parse internal doc (PDF, Markdown, HTML) and generate instruction-response pairs.'''
    validate_file_path(doc_path)
    logger.info(f'Processing document: {doc_path}')

    # Extract text using unstructured (handles PDF, MD, HTML)
    try:
        elements = partition.partition(str(doc_path))
    except Exception as e:
        logger.error(f'Failed to parse {doc_path}: {e}')
        return []

    full_text = '
'.join([el.text for el in elements if el.text.strip()])
    if not full_text.strip():
        logger.warning(f'No text extracted from {doc_path}')
        return []

    # Generate synthetic Q&A pairs using a small teacher model (optional, we used GPT-4o for this step)
    # Note: For reproducibility, we include pre-generated Q&A pairs in our public dataset
    chunks = chunk_document(full_text)
    qa_pairs = []

    # Load pre-generated Q&A pairs (stored as JSONL alongside docs)
    qa_path = doc_path.parent / f'{doc_path.stem}_qa.jsonl'
    if qa_path.exists():
        with open(qa_path, 'r') as f:
            for line in f:
                try:
                    qa_pairs.append(json.loads(line))
                except json.JSONDecodeError as e:
                    logger.error(f'Invalid JSON in {qa_path}: {e}')
    else:
        logger.warning(f'No pre-generated Q&A found for {doc_path}, skipping')
        return []

    # Format each Q&A pair into Llama 3.1 instruction template
    formatted_data = []
    for pair in qa_pairs:
        # Chunk context to fit max chunk size
        context_chunks = chunk_document(pair['context'])
        for ctx_chunk in context_chunks:
            formatted_text = LLAMA_3_1_INSTRUCTION_TEMPLATE.format(
                context=ctx_chunk,
                query=pair['query'],
                answer=pair['answer']
            )
            formatted_data.append({
                'text': formatted_text,
                'source_doc': str(doc_path),
                'query': pair['query']
            })

    logger.info(f'Generated {len(formatted_data)} samples from {doc_path}')
    return formatted_data

def main():
    # Configuration
    DOCS_DIR = Path('./internal_docs')
    OUTPUT_DATASET_PATH.mkdir(exist_ok=True)

    # Validate docs directory
    if not DOCS_DIR.exists():
        raise FileNotFoundError(f'Docs directory {DOCS_DIR} not found')

    all_samples = []
    # Process all supported document types
    supported_extensions = ['.pdf', '.md', '.html', '.txt']
    for ext in supported_extensions:
        for doc_path in DOCS_DIR.glob(f'**/*{ext}'):
            samples = process_internal_doc(doc_path)
            all_samples.extend(samples)

    if not all_samples:
        logger.error('No samples generated from any documents')
        return

    # Split into train/validation/test (80/10/10)
    np.random.seed(42)
    np.random.shuffle(all_samples)
    train_split = int(0.8 * len(all_samples))
    val_split = int(0.9 * len(all_samples))

    train_data = all_samples[:train_split]
    val_data = all_samples[train_split:val_split]
    test_data = all_samples[val_split:]

    # Save as Hugging Face dataset
    dataset = DatasetDict({
        'train': Dataset.from_list(train_data),
        'validation': Dataset.from_list(val_data),
        'test': Dataset.from_list(test_data)
    })

    dataset.save_to_disk(OUTPUT_DATASET_PATH)
    logger.info(f'Saved dataset with {len(train_data)} train, {len(val_data)} val, {len(test_data)} test samples to {OUTPUT_DATASET_PATH}')

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        logger.critical(f'Data preparation failed: {e}', exc_info=True)
        raise

import os
import torch
import logging
from pathlib import Path
from unsloth import FastLanguageModel, is_bfloat16_supported
from datasets import load_from_disk
from trl import SFTTrainer
from transformers import TrainingArguments, TextStreamer
import json

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.FileHandler('fine_tuning.log'), logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# Configuration
BASE_MODEL_ID = 'meta-llama/Meta-Llama-3.1-8B-Instruct'
DATASET_PATH = Path('./fine_tuning_data')
OUTPUT_MODEL_PATH = Path('./fine_tuned_llama3.1_8b_doc_assistant')
MAX_SEQ_LENGTH = 2048  # Llama 3.1 supports up to 128k, but we use 2k for cost efficiency
LOAD_IN_4BIT = True  # Use 4-bit quantization for lower memory usage
DTYPE = torch.bfloat16 if is_bfloat16_supported() else torch.float16

def validate_environment():
    '''Check for required GPU and dependencies before starting training.'''
    if not torch.cuda.is_available():
        raise RuntimeError('CUDA is not available. Fine-tuning requires an NVIDIA GPU.')
    gpu_mem = torch.cuda.get_device_properties(0).total_mem
    logger.info(f'Detected GPU: {torch.cuda.get_device_name(0)} with {gpu_mem / 1e9:.2f}GB VRAM')
    if LOAD_IN_4BIT and gpu_mem < 16e9:
        logger.warning('4-bit quantization recommended for GPUs with <24GB VRAM. Training may be slow.')
    if not DATASET_PATH.exists():
        raise FileNotFoundError(f'Dataset path {DATASET_PATH} not found. Run data_prep.py first.')

def load_model_and_tokenizer():
    '''Load Llama 3.1 8B with Unsloth optimizations.'''
    logger.info(f'Loading base model: {BASE_MODEL_ID}')
    try:
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=BASE_MODEL_ID,
            max_seq_length=MAX_SEQ_LENGTH,
            dtype=DTYPE,
            load_in_4bit=LOAD_IN_4BIT,
            token=os.getenv('HF_TOKEN'),  # Requires Hugging Face token for gated model
        )
    except Exception as e:
        logger.error(f'Failed to load model: {e}')
        raise

    # Add LoRA adapters for parameter-efficient fine-tuning
    model = FastLanguageModel.get_peft_model(
        model,
        r=16,  # LoRA rank
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],
        bias='none',
        use_gradient_checkpointing='unsloth',  # Unsloth's optimized gradient checkpointing
        random_state=42,
    )

    # Set padding token to EOS if not set
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = 'right'

    logger.info(f'Model loaded with {model.num_parameters():,} total parameters, {model.num_parameters(only_trainable=True):,} trainable')
    return model, tokenizer

def load_and_prepare_dataset(tokenizer):
    '''Load and tokenize the fine-tuning dataset.'''
    logger.info(f'Loading dataset from {DATASET_PATH}')
    dataset = load_from_disk(DATASET_PATH)

    def tokenize_function(examples):
        return tokenizer(
            examples['text'],
            truncation=True,
            max_length=MAX_SEQ_LENGTH,
            padding='max_length',
            return_tensors='pt'
        )

    # Tokenize all splits
    tokenized_dataset = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=dataset['train'].column_names
    )

    logger.info(f'Tokenized dataset sizes: Train {len(tokenized_dataset['train'])}, Val {len(tokenized_dataset['validation'])}')
    return tokenized_dataset

def train_model(model, tokenizer, tokenized_dataset):
    '''Run supervised fine-tuning with SFTTrainer.'''
    logger.info('Starting fine-tuning')
    training_args = TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_ratio=0.1,
        num_train_epochs=3,
        learning_rate=2e-4,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=10,
        save_strategy='epoch',
        evaluation_strategy='epoch',
        save_total_limit=2,
        load_best_model_at_end=True,
        metric_for_best_model='eval_loss',
        greater_is_better=False,
        output_dir=OUTPUT_MODEL_PATH / 'checkpoints',
        report_to='none',  # Disable wandb/tensorboard for simplicity
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=tokenized_dataset['train'],
        eval_dataset=tokenized_dataset['validation'],
        dataset_text_field='text',
        max_seq_length=MAX_SEQ_LENGTH,
        args=training_args,
    )

    try:
        trainer.train()
    except Exception as e:
        logger.error(f'Training failed: {e}', exc_info=True)
        raise

    # Save final model and tokenizer
    model.save_pretrained(OUTPUT_MODEL_PATH)
    tokenizer.save_pretrained(OUTPUT_MODEL_PATH)
    logger.info(f'Model saved to {OUTPUT_MODEL_PATH}')

    # Save training metrics
    metrics = trainer.evaluate(tokenized_dataset['test'])
    with open(OUTPUT_MODEL_PATH / 'test_metrics.json', 'w') as f:
        json.dump(metrics, f, indent=2)
    logger.info(f'Test metrics: {metrics}')

def main():
    validate_environment()
    model, tokenizer = load_model_and_tokenizer()
    tokenized_dataset = load_and_prepare_dataset(tokenizer)
    train_model(model, tokenizer, tokenized_dataset)

    # Test inference after training
    logger.info('Testing inference with fine-tuned model')
    FastLanguageModel.for_inference(model)
    streamer = TextStreamer(tokenizer)
    prompt = '''<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a technical documentation assistant for [Company Name Redacted] engineering team. Answer queries using only the provided context. If the answer is not in the context, say "I don't have enough information to answer that."<|eot_id|><|start_header_id|>user<|end_header_id|>
Context: Our CI pipeline uses GitHub Actions with a 10-minute timeout for unit tests. If tests exceed this, the job is cancelled.

Query: What happens if unit tests take longer than 10 minutes in our CI pipeline?<|eot_id|><|start_header_id|>assistant<|end_header_id|>'''
    inputs = tokenizer([prompt], return_tensors='pt').to('cuda')
    model.generate(**inputs, streamer=streamer, max_new_tokens=128)

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        logger.critical(f'Fine-tuning pipeline failed: {e}', exc_info=True)
        raise

import os
import time
import json
import logging
from pathlib import Path
from typing import List, Dict, Any, Tuple
import torch
from unsloth import FastLanguageModel
from transformers import TextStreamer
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank
from datasets import load_from_disk
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.FileHandler('benchmark.log'), logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# Configuration
FINE_TUNED_MODEL_PATH = Path('./fine_tuned_llama3.1_8b_doc_assistant')
RAG_VECTOR_STORE_PATH = Path('./rag_vector_store')
DATASET_PATH = Path('./fine_tuning_data')
BENCHMARK_OUTPUT_PATH = Path('./benchmark_results.json')
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
COHERE_API_KEY = os.getenv('COHERE_API_KEY')

# Validate environment variables
if not OPENAI_API_KEY:
    raise ValueError('OPENAI_API_KEY environment variable not set')
if not COHERE_API_KEY:
    raise ValueError('COHERE_API_KEY environment variable not set')

def load_fine_tuned_model():
    '''Load the fine-tuned Llama 3.1 model for inference.'''
    logger.info(f'Loading fine-tuned model from {FINE_TUNED_MODEL_PATH}')
    try:
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=str(FINE_TUNED_MODEL_PATH),
            max_seq_length=2048,
            dtype=torch.bfloat16,
            load_in_4bit=True,
        )
        FastLanguageModel.for_inference(model)
        return model, tokenizer
    except Exception as e:
        logger.error(f'Failed to load fine-tuned model: {e}')
        raise

def load_rag_pipeline():
    '''Initialize RAG 2.0 pipeline with FAISS, OpenAI embeddings, Cohere reranking.'''
    logger.info('Loading RAG 2.0 pipeline')
    if not RAG_VECTOR_STORE_PATH.exists():
        raise FileNotFoundError(f'RAG vector store not found at {RAG_VECTOR_STORE_PATH}. Run rag_index.py first.')

    # Load vector store
    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
    vector_store = FAISS.load_local(
        str(RAG_VECTOR_STORE_PATH),
        embeddings,
        allow_dangerous_deserialization=True  # Only use for trusted local indexes
    )

    # Base retriever (top 20 chunks)
    base_retriever = vector_store.as_retriever(search_kwargs={'k': 20})

    # Reranker (Cohere rerank, top 5 chunks)
    compressor = CohereRerank(cohere_api_key=COHERE_API_KEY, top_n=5)
    retriever = ContextualCompressionRetriever(
        base_compressor=compressor,
        base_retriever=base_retriever
    )

    # LLM for RAG (GPT-4o, same as we used in original RAG pipeline)
    llm = ChatOpenAI(model='gpt-4o', openai_api_key=OPENAI_API_KEY, temperature=0)

    return retriever, llm

def format_rag_prompt(context_docs: List[str], query: str) -> str:
    '''Format prompt for RAG pipeline with retrieved context.'''
    context = '

'.join([doc.page_content for doc in context_docs])
    return f'''You are a technical documentation assistant for [Company Name Redacted] engineering team. Answer queries using only the provided context. If the answer is not in the context, say "I don't have enough information to answer that."

Context:
{context}

Query: {query}

Answer:'''

def run_fine_tuned_inference(model, tokenizer, query: str, context: str = None) -> Tuple[str, float]:
    '''Run inference with fine-tuned Llama 3.1. Returns (answer, latency).'''
    if context is None:
        # For fine-tuned model, context is embedded in the training data, so we don't pass it at inference
        prompt = f'''<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a technical documentation assistant for [Company Name Redacted] engineering team. Answer queries using only your training data. If you don't know the answer, say "I don't have enough information to answer that."<|eot_id|><|start_header_id|>user<|end_header_id|>
Query: {query}<|eot_id|><|start_header_id|>assistant<|end_header_id|>'''
    else:
        prompt = f'''<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a technical documentation assistant for [Company Name Redacted] engineering team. Answer queries using only the provided context. If the answer is not in the context, say "I don't have enough information to answer that."<|eot_id|><|start_header_id|>user<|end_header_id|>
Context: {context}

Query: {query}<|eot_id|><|start_header_id|>assistant<|end_header_id|>'''

    inputs = tokenizer([prompt], return_tensors='pt').to('cuda')
    start_time = time.time()
    outputs = model.generate(**inputs, max_new_tokens=256, use_cache=True)
    latency = time.time() - start_time
    answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0].split('<|start_header_id|>assistant<|end_header_id|>')[-1].strip()
    return answer, latency

def run_rag_inference(retriever, llm, query: str) -> Tuple[str, float, List[str]]:
    '''Run inference with RAG 2.0 pipeline. Returns (answer, latency, retrieved_contexts).'''
    start_time = time.time()
    # Retrieve and rerank context
    docs = retriever.invoke(query)
    context = '

'.join([doc.page_content for doc in docs])
    # Generate answer
    prompt = format_rag_prompt(docs, query)
    answer = llm.invoke(prompt).content
    latency = time.time() - start_time
    return answer, latency, [doc.page_content for doc in docs]

def evaluate_results(predictions: List[str], ground_truths: List[str]) -> Dict[str, float]:
    '''Evaluate accuracy using exact match and F1 score (token-level).'''
    # Simple exact match (normalize whitespace)
    exact_match = []
    for pred, gt in zip(predictions, ground_truths):
        pred_norm = ' '.join(pred.lower().split())
        gt_norm = ' '.join(gt.lower().split())
        exact_match.append(1 if pred_norm == gt_norm else 0)

    # Token-level F1 (simplified)
    f1_scores = []
    for pred, gt in zip(predictions, ground_truths):
        pred_tokens = pred.lower().split()
        gt_tokens = gt.lower().split()
        common = set(pred_tokens) & set(gt_tokens)
        if not common:
            f1_scores.append(0.0)
            continue
        precision = len(common) / len(pred_tokens) if len(pred_tokens) > 0 else 0
        recall = len(common) / len(gt_tokens) if len(gt_tokens) > 0 else 0
        if precision + recall == 0:
            f1_scores.append(0.0)
            continue
        f1 = 2 * (precision * recall) / (precision + recall)
        f1_scores.append(f1)

    return {
        'exact_match': np.mean(exact_match),
        'token_f1': np.mean(f1_scores),
    }

def main():
    # Load models and pipelines
    ft_model, ft_tokenizer = load_fine_tuned_model()
    rag_retriever, rag_llm = load_rag_pipeline()

    # Load benchmark dataset (test split from fine-tuning data)
    dataset = load_from_disk(DATASET_PATH)
    test_data = dataset['test'].to_list()
    logger.info(f'Running benchmark on {len(test_data)} test samples')

    # Run inference for both models
    ft_results = []
    rag_results = []
    latencies_ft = []
    latencies_rag = []

    for sample in test_data:
        query = sample['query']
        ground_truth = sample['answer']

        # Fine-tuned model inference
        ft_answer, ft_latency = run_fine_tuned_inference(ft_model, ft_tokenizer, query)
        latencies_ft.append(ft_latency)
        ft_results.append({
            'query': query,
            'ground_truth': ground_truth,
            'prediction': ft_answer,
            'latency': ft_latency
        })

        # RAG 2.0 inference
        rag_answer, rag_latency, rag_contexts = run_rag_inference(rag_retriever, rag_llm, query)
        latencies_rag.append(rag_latency)
        rag_results.append({
            'query': query,
            'ground_truth': ground_truth,
            'prediction': rag_answer,
            'latency': rag_latency,
            'contexts': rag_contexts
        })

    # Evaluate
    ft_predictions = [r['prediction'] for r in ft_results]
    ft_ground_truths = [r['ground_truth'] for r in ft_results]
    ft_metrics = evaluate_results(ft_predictions, ft_ground_truths)
    ft_metrics['latency_p50'] = np.percentile(latencies_ft, 50)
    ft_metrics['latency_p99'] = np.percentile(latencies_ft, 99)

    rag_predictions = [r['prediction'] for r in rag_results]
    rag_ground_truths = [r['ground_truth'] for r in rag_results]
    rag_metrics = evaluate_results(rag_predictions, rag_ground_truths)
    rag_metrics['latency_p50'] = np.percentile(latencies_rag, 50)
    rag_metrics['latency_p99'] = np.percentile(latencies_rag, 99)

    # Save results
    results = {
        'fine_tuned_llama3.1': ft_metrics,
        'rag_2.0': rag_metrics,
        'test_samples': len(test_data),
        'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
    }

    with open(BENCHMARK_OUTPUT_PATH, 'w') as f:
        json.dump(results, f, indent=2)

    # Print summary
    logger.info('=== Benchmark Results ===')
    logger.info(f'Fine-Tuned Llama 3.1 8B: Exact Match {ft_metrics["exact_match"]:.2%}, F1 {ft_metrics["token_f1"]:.2%}, P99 Latency {ft_metrics["latency_p99"]:.2f}s')
    logger.info(f'RAG 2.0 (GPT-4o + Cohere Rerank): Exact Match {rag_metrics["exact_match"]:.2%}, F1 {rag_metrics["token_f1"]:.2%}, P99 Latency {rag_metrics["latency_p99"]:.2f}s')
    logger.info(f'Results saved to {BENCHMARK_OUTPUT_PATH}')

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        logger.critical(f'Benchmark failed: {e}', exc_info=True)
        raise

import json
from collections import defaultdict

def categorize_rag_errors(log_path: str) -> dict:
    '''Categorize RAG errors from JSONL log files.'''
    error_counts = defaultdict(int)
    with open(log_path, 'r') as f:
        for line in f:
            log = json.loads(line)
            if log['is_error']:
                # Check if context was retrieved
                if len(log['retrieved_contexts']) == 0:
                    error_counts['retrieval_failure'] += 1
                else:
                    # Check if ground truth is in retrieved context
                    gt_in_context = any(log['ground_truth'] in ctx for ctx in log['retrieved_contexts'])
                    if not gt_in_context:
                        error_counts['reranking_failure'] += 1
                    else:
                        error_counts['llm_failure'] += 1
    return dict(error_counts)

# Example usage
errors = categorize_rag_errors('./rag_logs.jsonl')
print(f'Retrieval: {errors.get("retrieval_failure", 0)}, Reranking: {errors.get("reranking_failure", 0)}, LLM: {errors.get("llm_failure", 0)}')

# axolotl_config.yml
base_model: meta-llama/Meta-Llama-3.1-8B-Instruct
model_type: LlamaForCausalLM
tokenizer_type: AutoTokenizer

load_in_4bit: true
adapter: lora
lora_r: 16
lora_alpha: 32
lora_dropout: 0.05
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj
  - gate_proj
  - up_proj
  - down_proj

num_epochs: 3
micro_batch_size: 2
gradient_accumulation_steps: 4
learning_rate: 2e-4
optimizer: adamw_bnb_8bit

python -m vllm.entrypoints.openai.api_server \
  --model ./fine_tuned_llama3.1_8b_doc_assistant \
  --tensor-parallel-size 1 \
  --gpu-memory-utilization 0.9 \
  --max-model-len 2048 \
  --port 8000