AI-Powered Document Q&A System with Retrieval-Augmented Generation

This document presents a production-ready architecture for building scalable document Q&A systems using Retrieval-Augmented Generation, covering data ingestion, vector storage, query processing, and system optimization strategies for enterprise deployment.

1. Problem Statement & ML Task Definition

Clarifying questions

1. What is the expected scale of documents and Q&A pairs?
   • Are we dealing with thousands, millions, or billions of documents and queries?
2. What is the latency budget for Q&A responses?
   • Should responses be sub-50ms for real-time interaction or can we tolerate 200–500ms?
3. What is the expected query volume and QPS?
   • Is this a high-throughput system (e.g., 10K+ QPS) or more modest usage?
4. What is the primary user interaction model?
   • Is this a chatbot-style system, or a search interface where users type queries and get document snippets?
5. What is the success definition for Q&A accuracy?
   • Is it top-1 answer correctness, or a broader measure like relevance score?
6. What is the expected freshness of documents?
   • Do we need real-time updates or can we batch process every hour/day?
7. What is the desired level of explainability?
   • Do we need to show sources or just provide answers?

Business use case and success criteria

The system should enable users to ask natural language questions about a collection of documents (e.g., technical manuals, legal contracts, or knowledge bases) and receive accurate, relevant answers directly from the documents. This supports use cases like intelligent search, chatbots, and automated document summarization.

Success criteria:

• Top-1 Answer Accuracy: Percentage of queries where the system returns the correct answer from the document set.
• Relevance Score: Average relevance of top-K answers (e.g., NDCG@5).
• User Engagement: Time spent on retrieved content, click-through rate on answer snippets.
• Latency: p99 response time under 200ms for interactive use cases.
• Coverage: Percentage of queries that can be answered from the document set.

ML problem formulation

This is a retrieval-augmented generation (RAG) task, which combines:

• Information retrieval to identify relevant document chunks.
• Natural language generation (NLG) to produce human-readable answers.

Justification:
The system must not only understand the query but also find and synthesize relevant information from a potentially large document corpus. This is fundamentally different from classification or ranking, as the output is generated from a combination of retrieved context and query understanding.

ML evaluation metrics

Offline Metrics:

• NDCG@K: Normalized Discounted Cumulative Gain to measure ranking quality.
• Precision@K: Proportion of relevant documents in top-K retrieved results.
• BLEU Score: For generated answer quality (if ground truth answers are available).
• ROUGE Score: For text generation relevance.

Online Metrics:

• Click-Through Rate (CTR) on answer snippets.
• User Engagement Time: Average time spent reading the retrieved content.
• Conversion Rate: Percentage of queries that lead to user actions (e.g., downloads, further questions).
• Query Relevance Score: User feedback on answer relevance.

Scale and constraints

Back-of-envelope estimates:

• Training data size: 10M query-document pairs (e.g., 100k queries × 100 documents each).
• Number of features: 50k unique tokens in vocabulary (for embeddings).
• Embedding matrix memory footprint:
  • 50k tokens × 768 dimensions × 4 bytes = ~150MB (per model).
• Serving QPS math:
  • If serving 10K QPS with 200ms p99 latency, we need ~2000 concurrent inferences.
  • Assuming a model with ~50ms inference time, we need 40 GPU instances (assuming 1 GPU = 50ms latency).

Constraints:

• Latency: p99 response time < 200ms.
• QPS: 10K+ queries/sec for interactive use cases.
• Data Volume: 1M+ documents, each up to 50k tokens.
• Freshness: Documents updated every 24 hours; queries must be served with up-to-date content.

Non-ML baseline strategy

• Rule-based fallback: For queries that don’t match any document, return a default "I cannot answer that" message.
• Keyword search: For simple queries, use keyword matching or TF-IDF to find relevant documents.
• Cold start handling: For new users or queries with no prior context, return a generic "Please provide more details" message.
• Hybrid approach: Use keyword search + semantic search for better coverage and faster fallback.

# Example of a simple rule-based fallback logic in Python
def fallback_query(query, document_store):
    if not query or len(query) < 3:
        return "Query too short. Please provide more details."
    if not document_store.contains(query):
        return "No documents found matching your query."
    return None  # Let ML system handle it

2. Data Strategy

Training data sources

• Internal document corpus: Company knowledge base, technical documentation, internal wikis, product specs, API docs (primary source)
• External knowledge sources: Publicly available documentation, StackOverflow, GitHub repositories, technical blogs (to expand domain coverage)
• Synthetic data: Generated question-answer pairs from existing documents using rule-based templates and LLMs to increase diversity
• User interaction logs: Query logs, click-through data, feedback signals from previous Q&A sessions (for relevance signal)

Labeling pipeline design

• Human annotation: For high-quality training data, particularly for edge cases and domain-specific questions (e.g., 1000+ high-value QA pairs from subject matter experts)
• Weak supervision: Use heuristics and rule-based matching to label document sections with relevant questions (e.g., "find all mentions of X in this document")
• Programmatic labeling: Auto-generate QA pairs from document structure (headings, tables, code snippets) using templates and LLMs for initial coverage

Class imbalance handling strategy with justification

• Problem: Most documents contain only a few relevant questions, leading to highly imbalanced training data where some document sections are underrepresented
• Strategy:
  • Use stratified sampling during training to ensure representation from all document categories
  • Apply focal loss or weighted loss functions to reduce impact of majority classes during training
  • Oversample rare question-document pairs in training batches using data augmentation techniques (e.g., paraphrasing questions or combining document sections)
• Justification: This ensures model learns to handle rare but important query-document relationships without being dominated by frequent, generic queries

Negative sampling approach

• Hard negative sampling: Sample document sections that are semantically similar but don't contain the answer to the query (e.g., documents on related topics)
• Random negative sampling: Randomly sample document sections from the corpus that don't match query intent (to simulate real-world search behavior)
• Negative mining: Use model predictions to identify hard negatives from inference on held-out data (active learning approach)
• Strategy: Combine hard and random negatives in ratio 3:1 to maintain model robustness while ensuring sufficient diversity

Data freshness and update frequency requirements

• Freshness requirement: Daily updates for internal documentation, weekly for external sources (e.g., StackOverflow)
• Update frequency:
  • High-priority documents (product specs, API docs): hourly updates
  • General documentation: daily updates
  • External sources: weekly crawl and update cycle
• Back-of-envelope estimate:
  • 500k documents with average 2k tokens → ~1GB training data (compressed)
  • 50k new questions per day → 1M Q&A pairs daily (with sampling)

Privacy compliance (GDPR, CCPA)

• Data anonymization: Remove PII from documents and queries before training (e.g., replace names with placeholders, redact email addresses)
• Data retention policy:
  • Internal documents: retain for 2 years after last modification
  • User queries: retain for 30 days (with consent for analytics)
• Compliance strategy:
  • Implement data access controls and audit logs
  • Use differential privacy techniques for training on sensitive documents
  • Provide user opt-out mechanisms for query logging

Data quality validation and monitoring

• Validation checks:
  • Semantic coherence between question and document content (e.g., use sentence transformers to verify relevance)
  • Answer completeness (check that answers fully address the question)
  • Document freshness (track last modified date and flag outdated content)
• Monitoring:
  • Daily drift detection on query distribution (PSI > 0.1 triggers alert)
  • Weekly human review of top-50 most frequent queries for quality assessment
  • Monthly data quality score (based on human annotation and model confidence scores)

# Example data validation function for QA pair quality check
def validate_qa_pair(question: str, answer: str, document_content: str) -> bool:
    """Validate that QA pair is semantically coherent and relevant"""
    # Use sentence transformer to compute similarity between question and document
    from sentence_transformers import SentenceTransformer
    
    model = SentenceTransformer('all-MiniLM-L6-v2')
    
    # Compute embeddings
    question_emb = model.encode(question)
    doc_emb = model.encode(document_content[:500])  # Truncate for efficiency
    
    # Compute cosine similarity
    from sklearn.metrics.pairwise import cosine_similarity
    similarity = cosine_similarity([question_emb], [doc_emb])[0][0]
    
    # Threshold for relevance (adjust based on domain)
    return similarity > 0.3 and len(answer) > 10

3. Feature Engineering

Feature Taxonomy

Features are organized into four categories:

1. User Features: User behavior signals, preferences, and demographics
2. Item Features: Document metadata, content characteristics, and categorization
3. Context Features: Query context, session information, and environmental signals
4. Interaction Features: User-document engagement patterns and query-document matching

Embedding Design

Entity Embeddings:

• Document IDs, categories, and user IDs → 128-dimensional dense embeddings
• Trained via contrastive learning with positive/negative document pairs

Text Embeddings:

• Query and document text → 768-dimensional BERT embeddings (base model)
• Fine-tuned on document Q&A pairs for domain relevance

Graph Embeddings:

• User-document interaction graph → 256-dimensional Node2Vec embeddings
• Captures user preferences and document similarity patterns

Real-time vs Batch Feature Computation

Feature Type	Computation Mode	Reasoning
User behavior signals (clicks, time spent)	Real-time	Immediate feedback for relevance
Document metadata (category, timestamp)	Batch	Static data, computed once per document update
Query embedding	Real-time	Dynamic query input requires fresh embeddings
User preferences (historical interactions)	Batch	Updated daily, cached for serving

Feature Store Design

Online Store:

• Low-latency Redis-based store for real-time feature serving
• TTL: 1 hour for user features, 24 hours for document features

Offline Store:

• Parquet-based data lake for training and batch processing
• Partitioned by date, user_id, and document_id

Feature Freshness and TTL

Feature	Source	Update Frequency	Store Type
Query text	User input	Real-time	Online
Document title	Document ingestion	Daily	Offline
User click history	User activity logs	Hourly	Online
Document category	Metadata extraction	Daily	Offline
User preference score	Training model	Daily	Online

Cross Features and Feature Interactions

Cross features are engineered to capture complex user-document relationships:

1. Query-Document Similarity Score: Cosine similarity between query and document embeddings
2. User-Document Engagement: User's historical interaction count with similar documents
3. Temporal Decay Features: Time-weighted relevance scores using exponential decay (λ=0.1)

Feature Importance and Selection

Methodology:

• SHAP values from trained XGBoost model to rank feature importance
• Permutation importance during cross-validation for robustness

Selection Strategy:

• Top 200 features retained based on cumulative importance > 95%
• Regular pruning to maintain model efficiency

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

class FeatureExtractor(nn.Module):
    def __init__(self, model_name="bert-base-uncased"):
        super().__init__()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.bert_model = AutoModel.from_pretrained(model_name)
        
    def forward(self, query_tokens, document_tokens):
        # Query embedding
        query_output = self.bert_model(**query_tokens)
        query_embedding = query_output.last_hidden_state.mean(dim=1)  # [batch_size, 768]
        
        # Document embedding  
        doc_output = self.bert_model(**document_tokens)
        doc_embedding = doc_output.last_hidden_state.mean(dim=1)  # [batch_size, 768]
        
        # Cosine similarity for cross-feature
        query_norm = torch.norm(query_embedding, dim=1, keepdim=True)
        doc_norm = torch.norm(doc_embedding, dim=1, keepdim=True)
        similarity = torch.sum(query_embedding * doc_embedding, dim=1) / (query_norm * doc_norm)
        
        return {
            "query_embedding": query_embedding,
            "doc_embedding": doc_embedding,
            "similarity_score": similarity
        }

4. Model Architecture

Model selection with justification

For a document Q&A system, I'll select a two-tower transformer architecture with the following rationale:

1. Retriever-Reader separation: Separating document retrieval from answer generation enables better scalability and allows independent optimization of each component
2. Transformer advantages: Transformers handle long-context documents well and can capture complex semantic relationships between queries and document passages
3. Scalability: Two-tower design allows for efficient retrieval with smaller models while maintaining strong answer generation capabilities

Model Stack:

• Retriever: Dense passage retriever (e.g., ColBERT or BM25 hybrid) - 100-200M parameters
• Reader: Extractive QA model (e.g., BERT-based or T5) - 300M-1B parameters
• Optional: Generative reader (e.g., LLM fine-tuned on QA data) - 10B+ parameters

Architecture diagram

graph TD
    A[Query Input] --> B[Query Encoder]
    C[Document Passages] --> D[Passage Encoder]
    B --> E[Query Embedding]
    D --> F[Passage Embeddings]
    E --> G{Retrieval}
    F --> G
    G --> H[Top-K Documents]
    H --> I[Answer Generation]
    I --> J[Answer Output]

Two-tower / multi-task / ensemble design

Two-Tower Architecture:

1. Query Encoder: BERT-base or similar transformer encoder
2. Passage Encoder: Same architecture as query encoder, but with different initialization

Multi-Task Learning Considerations:

• Joint training of retrieval and answer generation (if using same model)
• Separate training with cross-modal loss functions
• Ensemble of different architectures for robustness

Pre-trained model selection and fine-tuning strategy

Pre-trained Models:

• Query Encoder: mBERT (multilingual BERT) or Sentence-BERT for cross-lingual support
• Passage Encoder: Same as query encoder for consistency
• Answer Generator: BERT-QA or T5 for extractive generation, or LLaMA/Phi for generative

Fine-tuning Strategy:

1. Initial pre-training: Use sentence-transformers or mBERT with contrastive loss for retrieval
2. Fine-tuning:
   • Query encoder: Fine-tune on QA dataset (SQuAD, Natural Questions)
   • Passage encoder: Fine-tune with passage-level contrastive loss
   • Answer generator: Fine-tune on QA dataset with answer span extraction

Model complexity vs latency trade-off analysis

Model	Parameters	Latency (p99)	Memory Usage	Use Case
BERT-base (Retriever)	110M	50ms	2GB	Fast retrieval
BERT-base (Reader)	110M	100ms	2GB	Extractive QA
T5-small (Generative)	60M	200ms	1.5GB	Fast generative
LLaMA-7B (Generative)	7B	1500ms	20GB	High-quality generative

Offline evaluation baseline framing

Baseline Models:

1. BM25 + TF-IDF (traditional retrieval)
2. BERT-based extractive QA with no passage reranking
3. TF-IDF + rule-based answer extraction

Expected Performance:

• BM25 baseline: ~0.65 NDCG@10, ~0.45 F1
• BERT baseline: ~0.75 NDCG@10, ~0.65 F1
• Proposed model: ~0.85 NDCG@10, ~0.75 F1

Core model code snippet

import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer
from torch.nn.functional import cosine_similarity

class DocumentQAModel(nn.Module):
    def __init__(self, model_name='bert-base-uncased', dropout=0.1):
        super().__init__()
        self.query_encoder = BertModel.from_pretrained(model_name)
        self.passage_encoder = BertModel.from_pretrained(model_name)
        self.answer_generator = BertModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(dropout)
        
        # For answer generation, we'll use the encoder-decoder structure
        self.qa_head = nn.Linear(self.query_encoder.config.hidden_size, 2)  # start/end logits
        
    def forward(
        self,
        query_input_ids,
        query_attention_mask,
        passage_input_ids=None,
        passage_attention_mask=None,
        answer_start_positions=None,
        answer_end_positions=None
    ):
        # Query encoding (batch_size, query_seq_len)
        query_outputs = self.query_encoder(
            input_ids=query_input_ids,
            attention_mask=query_attention_mask
        )
        
        # Get query embeddings (batch_size, hidden_size)
        query_embeddings = query_outputs.pooler_output  # [batch_size, hidden_size]
        
        if passage_input_ids is not None:
            # Passage encoding (batch_size, passage_seq_len)
            passage_outputs = self.passage_encoder(
                input_ids=passage_input_ids,
                attention_mask=passage_attention_mask
            )
            
            # Get passage embeddings (batch_size, hidden_size)
            passage_embeddings = passage_outputs.pooler_output  # [batch_size, hidden_size]
            
            # Compute similarity for retrieval
            similarities = cosine_similarity(
                query_embeddings.unsqueeze(1), 
                passage_embeddings.unsqueeze(1)
            ).squeeze()
            
            return {
                'retrieval_scores': similarities,
                'query_embeddings': query_embeddings,
                'passage_embeddings': passage_embeddings
            }
        
        # For answer generation (extractive QA)
        if answer_start_positions is not None:
            # Use query and passage together for answer generation
            combined_input = torch.cat([query_input_ids, passage_input_ids], dim=1)
            combined_mask = torch.cat([query_attention_mask, passage_attention_mask], dim=1)
            
            # Forward through answer generator
            answer_outputs = self.answer_generator(
                input_ids=combined_input,
                attention_mask=combined_mask
            )
            
            # Extract answer start/end logits
            sequence_output = answer_outputs.last_hidden_state
            qa_logits = self.qa_head(sequence_output)
            
            start_logits, end_logits = qa_logits.split(1, dim=-1)
            start_logits = start_logits.squeeze(-1).contiguous()
            end_logits = end_logits.squeeze(-1).contiguous()
            
            # Compute loss if labels provided
            total_loss = None
            if answer_start_positions is not None:
                # Compute loss using cross-entropy for start/end positions
                loss_fct = nn.CrossEntropyLoss()
                start_loss = loss_fct(start_logits, answer_start_positions)
                end_loss = loss_fct(end_logits, answer_end_positions)
                total_loss = (start_loss + end_loss) / 2
            
            return {
                'loss': total_loss,
                'start_logits': start_logits,
                'end_logits': end_logits
            }
        
        # Default return for inference
        return {
            'query_embeddings': query_embeddings,
            'passage_embeddings': passage_embeddings if 'passage_embeddings' in locals() else None
        }

# Example usage for scoring function
def score_query_passage(query_embedding, passage_embedding):
    """Compute similarity score between query and passage"""
    return cosine_similarity(
        query_embedding.unsqueeze(0), 
        passage_embedding.unsqueeze(0)
    ).item()

# Example usage for answer generation
def generate_answer(model, query_tokens, passage_tokens):
    """Generate answer from query and passage"""
    # Tokenize inputs
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    
    # Prepare inputs
    query_input = tokenizer(
        query_tokens, 
        return_tensors='pt', 
        padding=True, 
        truncation=True
    )
    
    passage_input = tokenizer(
        passage_tokens, 
        return_tensors='pt', 
        padding=True, 
        truncation=True
    )
    
    # Forward pass
    outputs = model(
        query_input_ids=query_input['input_ids'],
        query_attention_mask=query_input['attention_mask'],
        passage_input_ids=passage_input['input_ids'],
        passage_attention_mask=passage_input['attention_mask']
    )
    
    # Extract answer (simplified for example)
    return outputs

This implementation provides a concrete foundation for building a document Q&A system with proper embedding computation, retrieval scoring, and answer generation capabilities. The model supports both extractive and generative approaches while maintaining the flexibility for future extensions.

5. Serving & Inference Architecture

System Diagram

graph TD
    A[User Query] --> B[Query Preprocessing]
    B --> C[Embedding Model]
    C --> D[Vector Store]
    D --> E[Retrieval Engine]
    E --> F[Contextual Prompt]
    F --> G[LLM Inference]
    G --> H[Response Generation]
    H --> I[Response Formatter]
    I --> J[User Response]

    subgraph "Online Serving Layer"
        B
        C
        E
        F
        G
        H
        I
    end

    subgraph "Storage Layer"
        D
    end

Online vs Batch Serving Decision

Online serving is required for interactive document Q&A where users expect sub-second response times. Batch processing would introduce unacceptable latency for real-time interactions. The system must support:

• P99 latency target: < 2 seconds for end-to-end response
• QPS capacity: 10,000+ concurrent queries
• Real-time context: Query-dependent response generation

Two-Stage Pipeline

1. Retrieval Stage:

   • Vector similarity search (faiss/Weaviate)
   • Top-K document retrieval (K=50)
   • Context windowing and formatting

2. Scoring/Generation Stage:

   • Prompt construction with retrieved context
   • LLM inference (vLLM or HuggingFace)
   • Response generation and validation

Serving Infrastructure

Primary Inference Engine: vLLM for LLM serving due to:

• High throughput (1000+ tokens/sec)
• Low latency (< 50ms for prompt processing)
• Efficient KV cache management
• Support for multi-GPU deployment

Model Serving Stack:

# vLLM serving configuration example
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3.2-1B-Instruct",
    tensor_parallel_size=4,
    max_model_len=8192,
    gpu_memory_utilization=0.8
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=512
)

Latency Budget Breakdown

Stage	p50 Latency	p99 Latency	Budget
Query Preprocessing	5ms	10ms	15ms
Embedding Generation	20ms	35ms	50ms
Vector Search	15ms	40ms	50ms
Prompt Construction	10ms	25ms	35ms
LLM Inference	100ms	200ms	250ms
Response Formatting	5ms	15ms	20ms
Total	155ms	360ms	420ms

Model Compression Strategy

• Quantization: 4-bit GPTQ quantization for LLMs (50% size reduction, 5-10% accuracy drop)
• Distillation: Knowledge distillation to smaller model variants (e.g., from Llama-3.2-1B to 0.5B)
• Pruning: Structured pruning for embedding layers (20% reduction in parameters)

Caching Strategy

Prediction Cache:

• Response cache with TTL of 1 hour for identical queries
• Cache key: hash(query + document_id)

Embedding Cache:

• Pre-computed embeddings for frequently accessed documents
• Cache invalidation on document updates

KV Cache:

• vLLM's built-in KV cache for LLM attention computation
• Memory optimized with sliding window attention

Auto-scaling and Capacity Planning

Horizontal Scaling:

• Auto-scaling groups for embedding models (2-4 instances)
• Load balancer with weighted round-robin distribution
• Kubernetes deployment with HPA (Horizontal Pod Autoscaler)

Capacity Planning:

• Target: 500 QPS per GPU (assuming 8x A100s)
• Expected peak: 2,500 QPS across 5 GPU nodes
• Cold start handling with warm-up routines

Fallback Strategy:

• Rule-based response for failed LLM inference
• Cached responses for frequently asked questions
• Graceful degradation to summary-based answers

6. Evaluation Framework

Offline Evaluation Methodology

For offline evaluation, we'll use temporal cross-validation to ensure our model generalizes well across time. We'll split the dataset chronologically, using 80% of documents for training and 20% for validation. This approach addresses the temporal nature of document Q&A where newer documents may have different content patterns than older ones. We'll also implement k-fold temporal splits (k=5) to ensure robustness across different time windows.

Slice-Based Evaluation

We'll evaluate model performance across critical slices to identify potential biases:

• User segments: New vs. returning users, power users vs. casual users
• Document types: Technical documentation, marketing content, support articles
• Geographic regions: Different language support, regional content variations
• Content freshness: Recent vs. older documents (e.g., 30-day vs. 90+ day old content)
• Query complexity: Simple vs. multi-hop questions

A/B Test Design

Randomization unit: User-query pairs to avoid query leakage between test groups.

Sample size calculation:

For 5% absolute difference in MRR@1 with 95% confidence and 80% power:
n = (Zα/2 + Zβ)² × (p1(1-p1) + p2(1-p2)) / (p1 - p2)²
Assuming baseline MRR@1 = 0.6, target = 0.65, n ≈ 1,200 queries per group

Test duration: 2 weeks for statistical significance with daily traffic of ~10k queries.

Primary metric: MRR@1 (Mean Reciprocal Rank) to measure ranking quality of top document matches.

Secondary metrics:

• Precision@1, Precision@3, NDCG@3
• User engagement (click-through rate on top result)
• Query success rate (user confirms answer is correct)

Guardrail Metrics and Blocking Conditions

Guardrail metrics:

• Response time: p99 > 500ms → immediate alert
• Query success rate: < 60% → halt traffic to new model version
• Hallucination rate: > 5% of responses contain fabricated content → immediate rollback

Blocking conditions:

• If MRR@1 drops by >10% compared to baseline
• If response latency exceeds 100ms p99 threshold
• If user feedback indicates >20% of queries are unsatisfactory

Shadow Mode Rollout Plan

1. Week 1: Route 5% of traffic through shadow mode
2. Week 2: Increase to 10%, monitor performance parity
3. Week 3: Route 25% of traffic, begin A/B testing with control group
4. Week 4: Full rollout to 100% if A/B test shows significant improvement

Long-Term Metric Measurement

Delayed conversion tracking:

• 1-day, 7-day, and 30-day user engagement with retrieved documents
• Document usage patterns (time spent reading, sections accessed)
• User feedback collection through post-query surveys

Engagement decay analysis:

• Track how quickly user engagement drops after initial query
• Monitor document re-usage patterns to detect content obsolescence

User retention impact:

• Measure if Q&A system improves overall document engagement
• Track user journey from query to full document consumption

# Example offline evaluation script for MRR calculation
import numpy as np
from typing import List, Tuple

def calculate_mrr(predictions: List[List[str]], ground_truths: List[str]) -> float:
    """
    Calculate Mean Reciprocal Rank for document retrieval
    predictions: list of ranked document IDs per query
    ground_truths: list of correct document IDs per query
    """
    mrr_scores = []
    
    for pred_list, true_doc in zip(predictions, ground_truths):
        try:
            rank = pred_list.index(true_doc) + 1  # 1-indexed
            mrr_scores.append(1.0 / rank)
        except ValueError:
            # If ground truth not in predictions, treat as worst case
            mrr_scores.append(0.0)
    
    return np.mean(mrr_scores)

def calculate_precision_at_k(predictions: List[List[str]], ground_truths: List[str], k: int = 3) -> float:
    """
    Calculate Precision@K for document retrieval
    """
    precision_scores = []
    
    for pred_list, true_doc in zip(predictions, ground_truths):
        relevant_docs = set(pred_list[:k])
        if true_doc in relevant_docs:
            precision_scores.append(1.0)
        else:
            precision_scores.append(0.0)
    
    return np.mean(precision_scores)

# Example usage:
# predictions = [["doc1", "doc2", "doc3"], ["doc4", "doc5", "doc6"]]
# ground_truths = ["doc2", "doc4"]
# mrr = calculate_mrr(predictions, ground_truths)
# precision_3 = calculate_precision_at_k(predictions, ground_truths, k=3)

7. Training Pipeline

Training Data Pipeline

The training pipeline begins with data ingestion from multiple sources including internal document repositories, external knowledge bases, and synthetic data generation. The pipeline includes:

1. Data Ingestion: Raw documents are ingested via document stores (e.g., MongoDB, S3) with metadata extraction for filtering and categorization
2. Preprocessing: Documents are chunked into 512-token segments with overlap, cleaned for noise, and normalized
3. Query Generation: For each document chunk, synthetic queries are generated using rule-based templates and LLM-powered query expansion
4. Data Shuffling: Training data is shuffled using reservoir sampling to ensure diverse distribution across batches

class DocumentQATrainingPipeline:
    def __init__(self, chunk_size=512, overlap=64):
        self.chunk_size = chunk_size
        self.overlap = overlap
    
    def process_documents(self, documents):
        chunks = []
        for doc in documents:
            # Chunk document with overlap
            tokens = self.tokenize(doc.content)
            for i in range(0, len(tokens), chunk_size - overlap):
                chunk = tokens[i:i + chunk_size]
                chunks.append({
                    'chunk_id': f"{doc.id}_{i}",
                    'content': self.detokenize(chunk),
                    'source_doc_id': doc.id,
                    'metadata': doc.metadata
                })
        return chunks
    
    def generate_training_data(self, chunks):
        training_examples = []
        for chunk in chunks:
            # Generate synthetic queries
            queries = self.generate_queries(chunk['content'])
            for query in queries:
                training_examples.append({
                    'query': query,
                    'context': chunk['content'],
                    'answer': self.extract_answer(query, chunk['content'])
                })
        return training_examples

Distributed Training Strategy

For large-scale model training, we implement a hybrid approach:

1. Data Parallelism: Distributed across multiple GPUs with gradient synchronization via NCCL
2. Model Parallelism: For large transformer models, we split layers across GPUs using tensor parallelism
3. Pipeline Parallelism: For very large models, we pipeline operations across stages

Training infrastructure:

• 8x A100 GPUs (80GB VRAM each) for model training
• 2x NVIDIA DGX-2 nodes (16x A100) for larger model experiments
• Estimated training time: 48 hours for full model fine-tuning on 1M samples

Hyperparameter Tuning Approach

We use Bayesian optimization for efficient hyperparameter search:

Key hyperparameters:

• Learning rate (1e-5 to 1e-3)
• Batch size (16, 32, 64)
• Dropout rate (0.1 to 0.5)
• Number of attention heads (8, 12, 16)
• Model layers (12, 24)

from skopt import gp_minimize
from skopt.space import Real, Integer

def objective(params):
    lr, batch_size, dropout = params
    # Train model with given hyperparameters
    model = train_model(
        learning_rate=lr,
        batch_size=batch_size,
        dropout=dropout
    )
    # Return validation loss (minimize)
    return validate_model(model)

# Bayesian optimization search space
search_space = [
    Real(1e-5, 1e-3, name='learning_rate'),
    Integer(16, 64, name='batch_size'),
    Real(0.1, 0.5, name='dropout')
]

best_params = gp_minimize(
    func=objective,
    dimensions=search_space,
    n_calls=50,
    random_state=42
)

Experiment Tracking and Reproducibility

We use MLflow for experiment tracking:

1. Experiment Registration: Each training run is logged with parameters, metrics, and artifacts
2. Version Control: Model checkpoints are versioned using Git LFS for reproducibility
3. Artifact Storage: Training logs, model weights, and evaluation results stored in cloud storage

Model Versioning and Registry Strategy

• Versioning: Semantic versioning (v1.0.0, v1.1.0) with major/minor/patch tags
• Registry: Model registry with staging and production environments
• Rollback Capability: Automated rollback to previous stable versions on performance degradation

Continual / Online Learning Design

For continuous model improvement:

1. Feedback Loop: User interactions (clicks, corrections) are collected and used for retraining
2. Incremental Updates: Periodic model updates every 2 weeks using new data samples
3. Online Learning: For lightweight models, we implement online learning with adaptive learning rates

Training frequency:

• Full retraining: Weekly (for major model updates)
• Incremental updates: Daily (for feedback-driven improvements)
• Shadow training: Continuous (for A/B testing of new model versions)

The pipeline supports both batch and streaming data ingestion with automatic scaling based on data volume and system load.

8. Monitoring & Observability

Data Drift Detection

Data drift detection is critical for maintaining system performance as document content and user behavior evolve over time. We implement drift monitoring using:

• PSI (Population Stability Index) for categorical features like document categories, user segments, and query types
• KL Divergence for continuous features like embedding distances and relevance scores
• Statistical Tests (Kolmogorov-Smirnov test) for feature distribution shifts

import numpy as np
from scipy import stats
import pandas as pd

def calculate_psi(expected, actual, bins=10):
    """Calculate Population Stability Index between two distributions"""
    # Ensure both arrays have same bins
    _, bin_edges = np.histogram(expected, bins=bins)
    
    # Count occurrences in each bin
    expected_counts, _ = np.histogram(expected, bins=bin_edges)
    actual_counts, _ = np.histogram(actual, bins=bin_edges)
    
    # Avoid zero division
    expected_counts = np.clip(expected_counts, 1e-6, None)
    actual_counts = np.clip(actual_counts, 1e-6, None)
    
    psi = np.sum((actual_counts - expected_counts) * 
                 np.log(actual_counts / expected_counts))
    return psi

def detect_drift_features(current_features, reference_features, threshold=0.1):
    """Detect drift in feature distributions"""
    drift_results = {}
    
    for feature_name, current_data in current_features.items():
        if feature_name in reference_features:
            psi_value = calculate_psi(
                reference_features[feature_name], 
                current_data
            )
            
            drift_results[feature_name] = {
                'psi': psi_value,
                'drift_detected': psi_value > threshold,
                'severity': 'high' if psi_value > 0.2 else 
                           'medium' if psi_value > 0.1 else 'low'
            }
    
    return drift_results

Prediction Drift and Label Drift Monitoring

Prediction drift monitors shifts in model outputs, while label drift tracks changes in ground truth distributions:

• Prediction Drift: Monitor change in ranking scores, relevance probability distributions, and response times
• Label Drift: Track distribution shifts in human judgments or click-through labels

Model Performance Degradation Alerting

We implement a multi-tiered alerting system:

• Primary Metrics: NDCG@K, MRR, Precision@K (threshold: 5% degradation)
• Secondary Metrics: Response time (p99 > 200ms), cache hit rate (<85%)
• Alerting Cadence: Hourly checks with 15-minute alert windows

class ModelPerformanceMonitor:
    def __init__(self, metrics_config):
        self.metrics_config = metrics_config
        self.alert_thresholds = {
            'ndcg': 0.05,      # 5% degradation
            'precision_k': 0.03,
            'latency_p99': 0.20 # 20% increase
        }
    
    def check_degradation(self, current_metrics, baseline_metrics):
        alerts = []
        
        for metric_name, threshold in self.alert_thresholds.items():
            if metric_name in current_metrics and metric_name in baseline_metrics:
                current = current_metrics[metric_name]
                baseline = baseline_metrics[metric_name]
                
                if current < baseline * (1 - threshold):
                    alerts.append({
                        'metric': metric_name,
                        'current': current,
                        'baseline': baseline,
                        'change_percent': (current - baseline) / baseline * 100,
                        'severity': 'high' if abs(current - baseline) / baseline > 0.1 else 'medium'
                    })
        
        return alerts

Retraining Trigger Strategy

We implement a hybrid retraining trigger system:

• Scheduled Retraining: Weekly model updates using new document collections
• Drift-Based Triggers:
  • PSI > 0.15 for any feature → trigger retraining
  • NDCG drop > 3% over 2 weeks → trigger retraining
• Performance-Based Triggers:
  • Latency > 500ms for 3 consecutive hours → model refresh

Logging Design

Comprehensive logging includes:

• Prediction Logs: Query, document IDs, scores, response time, cache hits
• Feature Logs: User features, query embeddings, document embeddings
• Label Collection Pipeline: Human feedback, click data, relevance judgments

import json
from datetime import datetime

class PredictionLogger:
    def __init__(self, log_file_path):
        self.log_file = log_file_path
    
    def log_prediction(self, query, documents, scores, response_time, 
                      user_id=None, session_id=None):
        log_entry = {
            'timestamp': datetime.utcnow().isoformat(),
            'query': query,
            'documents': documents,
            'scores': scores,
            'response_time_ms': response_time,
            'user_id': user_id,
            'session_id': session_id,
            'cache_hit': False
        }
        
        with open(self.log_file, 'a') as f:
            f.write(json.dumps(log_entry) + '\n')

Feedback Loop for Continuous Improvement

We establish a feedback loop with:

• Human-in-the-Loop: Periodic manual evaluation of top-ranked documents
• Click-Through Analysis: Track user engagement with retrieved documents
• Relevance Feedback: Collect explicit relevance judgments from users
• Model Retraining Pipeline: Automated retraining when feedback thresholds exceeded

The monitoring system is designed to detect degradation within 24 hours and trigger retraining or alerting when performance drops below acceptable thresholds, ensuring system reliability and continuous improvement.

9. Responsible AI & Safety

Bias Analysis

• Training data audit: Conduct comprehensive audit of document corpus for demographic, geographic, and topic representation bias. Analyze distribution of content by domain (technical, medical, legal, etc.) to ensure balanced coverage.
• Slice-based fairness evaluation: Evaluate model performance across user segments (age groups, education levels, geographic regions) to detect potential disparities in answer quality or relevance.
• Content bias detection: Implement automated checks for gender, racial, and cultural bias in generated answers through toxicity classification models and bias detection libraries.

Fairness Metrics and Constraints

• Primary fairness metric: Equalized odds across different user demographics to ensure consistent answer quality regardless of user characteristics.
• Decision rationale: Equalized odds is preferred over demographic parity as it balances the trade-off between accuracy and fairness, particularly important in Q&A systems where answer correctness is paramount.
• Implementation approach: Monitor precision and recall for different slices and implement fairness constraints during model training via weighted loss functions or adversarial debiasing techniques.

Explainability Requirements and Approach

• Stakeholder requirements:
  • Business stakeholders: Need to understand which document sections contributed to answers for compliance and content validation
  • End users: Require transparency in answer generation process for trust building
• Explainability methods:
  • Attention visualization for transformer-based models to show document passage relevance
  • SHAP values for feature importance in retrieval components
  • Document-level confidence scores to indicate answer reliability
• Implementation: Integrate explainability modules that return passage relevance scores and attention weights alongside answers, with optional UI display for end-users.

LLM Safety Guardrails

• Content filtering:
  • Implement content safety classifiers to detect and filter out harmful, illegal, or inappropriate responses
  • Use toxicity detection models (e.g., Perspective API) to flag potentially unsafe content
• Output validation:
  • Validate answer facts against document sources to prevent hallucination
  • Implement confidence thresholds below which answers are rejected or flagged for human review
• Rate limiting:
  • Enforce rate limits to prevent abuse and ensure system stability
  • Implement request throttling based on user session patterns

Model Card Documentation Plan

• Content requirements:
  • Model architecture and training details (pre-trained model used, fine-tuning strategy)
  • Performance metrics across different data slices and domains
  • Known limitations (e.g., handling of ambiguous queries, domain-specific knowledge gaps)
  • Bias and fairness considerations with mitigation strategies
• Distribution:
  • Maintain model cards in version-controlled documentation system accessible to ML engineers and product teams
  • Include sample inputs/outputs with explanations for common use cases

Incident Response Plan for Model Failures

• Failure detection:
  • Monitor for sudden drops in answer quality, increased hallucination rates, or unexpected behavior patterns
  • Implement alerting for system-wide performance degradation (e.g., p99 latency > 500ms)
• Response protocol:
  • Immediate rollback to previous stable model version if accuracy drops below threshold (e.g., <85% F1 score)
  • Flag and log all failed queries for human review and potential retraining trigger
  • Notify relevant stakeholders (product, engineering, compliance) for incident triage and communication

# Example safety validation function for RAG system output
def validate_answer_safety(answer_text, retrieved_docs, model_config):
    """
    Validate answer for safety and factual correctness
    Returns: (is_safe, safety_reasons, confidence_score)
    """
    # Check for toxicity
    toxicity_score = detect_toxicity(answer_text)
    
    # Check for hallucination (answer not supported by documents)
    is_factual = check_factual_consistency(answer_text, retrieved_docs)
    
    # Apply confidence threshold
    confidence = calculate_answer_confidence(answer_text, retrieved_docs)
    
    safety_issues = []
    if toxicity_score > 0.8:
        safety_issues.append("toxic_content")
    
    if not is_factual:
        safety_issues.append("hallucination")
    
    is_safe = len(safety_issues) == 0
    return is_safe, safety_issues, confidence

def detect_toxicity(text):
    """Simple toxicity detection using pre-trained classifier"""
    # Placeholder for actual toxicity detection model call
    return 0.0 if len(text) < 10 else 0.2  # Mock implementation

def check_factual_consistency(answer, docs):
    """Check if answer is supported by retrieved documents"""
    # Simple keyword-based check for demonstration
    answer_words = set(answer.lower().split())
    doc_content = " ".join([doc.get('content', '') for doc in docs])
    doc_words = set(doc_content.lower().split())
    
    # If answer contains significant portion of document content, likely factual
    overlap = len(answer_words.intersection(doc_words))
    return overlap > (len(answer_words) * 0.3)

10. Scalability & Reliability

Bottleneck Analysis per Pipeline Stage

Data Ingestion Pipeline:

• Throughput: 100GB/hour processing capacity needed for enterprise-scale document ingestion
• Bottleneck: Vector embedding generation (100k documents/sec → 50k embeddings/sec per GPU)
• Back-of-envelope: With 8x A100 GPUs, can process ~50k docs/sec, assuming 200ms/doc embedding time

Retrieval Stage:

• Query Processing: 10k QPS target → 50ms p99 latency budget
• Vector Search: 1M document vectors, 768-dim embeddings → 3GB memory footprint per GPU
• Bottleneck: Multi-tenant query routing and concurrent retrieval across 1000+ document collections

Scoring & Ranking:

• Model Inference: 10k QPS → 5ms p99 latency per query for reranking (200ms total)
• Batching: 64 queries per batch for optimal GPU utilization
• Memory: 10GB VRAM per model instance (assuming 7B parameter model)

Scale-up Strategy

Horizontal Scaling of Serving:

• Query Processing: Auto-scale to 50+ replica instances for 10k QPS
• Vector Stores: Sharding across 100+ vector databases (Pinecone/Weaviate clusters)
• Feature Store: 100+ feature store shards for concurrent user query processing

Distributed Feature Computation:

• Real-time Features: Kafka-based streaming for 10k QPS feature updates
• Batch Feature Computation: Spark-based batch jobs with 100+ executors for large-scale feature engineering

Fault Tolerance

Model Serving Fallbacks:

• Primary: 3x redundancy with auto-failover to backup model instances
• Fallback: Rule-based query routing to cached responses when model fails
• Stale Model Serving: 5-minute grace period for model warm-up during failover

Stale Model Handling:

• Cache Invalidation: 15-minute TTL for model predictions
• Graceful Degradation: Fallback to TF-IDF retrieval when vector store is unavailable
• Rollback Strategy: 5-minute rollback window for model versions with A/B test validation

Cold Start Handling Strategy

New User Handling:

• Fallback: Default to popular document rankings for new users
• Personalization: Use demographic features and historical behavior patterns (if available)
• Grace Period: 24-hour cold start window with rule-based fallback

New Document Handling:

• Indexing Pipeline: 24-hour indexing window for new documents
• Fallback Retrieval: Use TF-IDF or BM25 as fallback until vector indexing completes
• Hybrid Retrieval: 70% vector + 30% keyword search for new documents

Multi-region Serving Strategy

Global Distribution:

• Primary Region: US-East for low-latency response (50ms p99)
• Secondary Regions: EU-West and AP-Southeast for regional redundancy
• Cross-Region Sync: 1-hour sync window for feature and model updates

Traffic Distribution:

• Load Balancing: 60% US-East, 20% EU-West, 20% AP-Southeast
• Failover: 30-second failover window for region-level outages
• Consistency: Eventual consistency model with 1-hour global sync window

Cross-Region Latency Budget:

• p99: 150ms for cross-region queries
• p50: 80ms for primary region queries
• Total Query Latency: 200ms p99 (including retrieval, reranking, and response generation)

Infrastructure Redundancy:

• Model Replicas: 3x per region with cross-region replication
• Vector Stores: Multi-region Pinecone clusters with 1-hour sync window
• Feature Store: 3x redundancy across regions with 15-minute TTL for stale features

Capacity Planning:

• QPS Growth: 2x annual growth rate projected for enterprise Q&A systems
• Scaling Factor: 10x capacity planning for peak traffic (20k QPS) with 50ms p99 target
• Resource Utilization: 70% GPU utilization target with 30% headroom for scaling

11. Open Questions, Risks & Trade-offs

Decisions requiring product / engineering / policy alignment

• User intent classification: Should we prioritize exact match queries or semantic understanding? This affects whether to use dense retrieval vs sparse retrieval approaches.
• Document source trustworthiness: How do we handle documents from untrusted sources (e.g., web scraping vs internal documentation)? This impacts data quality and safety requirements.
• Response generation strategy: Should we prefer exact quotes from source documents or paraphrased answers? This affects model training objectives and output quality.
• Multi-turn conversation support: Do we need to maintain context across multiple Q&A exchanges? This affects system architecture and memory management requirements.

Key ML risks

Risk	Impact	Likelihood	Mitigation
Hallucination	High - generates false information	Medium-High	Implement strict source citation, use retrieval-augmented generation with confidence thresholds, red teaming
Data drift in document corpus	High - outdated documents lead to incorrect answers	Medium	Implement document freshness monitoring, scheduled retraining with new content
Retrieval failure	Medium - no relevant documents found	High	Implement fallback to broader search, hybrid retrieval strategies
Query misunderstanding	Medium - wrong answer due to misinterpretation	Medium	Use query reformulation, multi-lingual support, and user feedback loops
Privacy leakage	High - sensitive information in documents	Medium	Implement document-level access controls, content filtering, and anonymization

Trade-offs made with justification

Trade-off	Justification
Accuracy vs Latency	Prioritize sub-500ms response time for user experience over marginal accuracy gains from deeper models
Freshness vs Cost	Accept some staleness in document embeddings to reduce computational cost of frequent retraining
Model Size vs Quality	Use smaller, efficient models with quantization to enable serving on edge devices while maintaining >90% accuracy
Custom vs Pre-trained Models	Leverage pre-trained transformers for faster time-to-market, but fine-tune on domain-specific data
Online vs Batch Training	Use batch training for stability and reproducibility, with online learning for rapid adaptation to new documents

Alternatives considered and rejected

• Traditional keyword-based search: Rejected due to poor semantic understanding and inability to handle nuanced queries
• End-to-end training without retrieval: Rejected due to scalability issues with large document collections and hallucination risks
• Pure LLM without RAG: Rejected due to poor factual accuracy and inability to provide source attribution
• Single-tower architecture: Rejected due to suboptimal retrieval quality compared to two-tower or transformer-based approaches

Technical debt acknowledged

• Hardcoded document types: Current system assumes all documents are in a single format, requiring future refactoring for multi-format support
• Limited query understanding: Basic intent classification may not scale to complex multi-hop queries without significant rework
• No cross-document reasoning: Current design doesn't support answering questions that require information from multiple documents
• Limited explainability: Model outputs lack clear attribution to specific document passages, making debugging difficult

12. Phased Implementation Roadmap

Phase 1: MVP (Weeks 1-6)

Objective: Establish baseline document Q&A system with basic RAG capabilities

Deliverables:

• Rule-based fallback system (keyword matching, document metadata search)
• Simple embedding model (e.g., sentence-transformers/all-MiniLM-L6-v2)
• Basic document ingestion pipeline (PDF, DOCX, TXT support)
• Single-vector store for document embeddings
• Simple query processing pipeline (query → embedding → similarity search → answer generation)

Dependencies:

• Document storage system (S3-compatible)
• Basic feature store for metadata
• Simple model serving infrastructure

Phase 2: Scale (Weeks 7-12)

Objective: Expand to production-grade system with distributed capabilities

Deliverables:

• Multi-vector store (semantic + keyword + graph embeddings)
• Distributed feature computation pipeline
• Batch processing for document embedding generation
• Multi-stage retrieval pipeline (candidate filtering → semantic ranking → re-ranking)
• A/B testing infrastructure for model variants
• Model versioning and registry system

Dependencies:

• Distributed training cluster (GPU/TPU)
• Feature store with online/offline capabilities
• ML experiment tracking system

Phase 3: Excellence (Weeks 13-18)

Objective: Advanced capabilities with continuous learning and real-time adaptation

Deliverables:

• Multi-task model architecture (retrieval + generation + relevance scoring)
• Real-time feature computation for user context and query dynamics
• Online learning pipeline for continuous model improvement
• Advanced reranking with cross-encoder models
• Multi-modal document support (images, tables)
• Comprehensive monitoring and alerting system

Dependencies:

• Real-time feature processing pipeline
• Advanced model serving infrastructure (vLLM, Triton)
• Comprehensive logging and observability stack

Critical Path Dependencies:

1. Feature Store - Must be ready before Phase 2 to support distributed training
2. Model Serving Infrastructure - Critical for end-to-end latency in Phase 2
3. Document Ingestion Pipeline - Foundation for all downstream systems
4. A/B Testing Framework - Required for Phase 2 and 3 validation

Risk Mitigation:

• Data Pipeline Bottlenecks: Pre-compute embeddings in Phase 1, optimize batch sizes
• Model Performance Degradation: Maintain rule-based fallback during Phase 2 rollout
• Scalability Issues: Design distributed components from Phase 1 to avoid rework

Success Metrics:

• Phase 1: <500ms query latency, >70% document retrieval accuracy
• Phase 2: <200ms query latency, >85% document retrieval accuracy, 3x throughput improvement
• Phase 3: <100ms query latency, >90% document retrieval accuracy, continuous model improvement