AI-Powered Code Completion System Design

Comprehensive system design for Copilot-like code completion engines with neural architecture, real-time inference, and multi-language support.

PublishedJanuary 27, 2026
Length6,875 words
Reading time35 min read

AI/ML System Design

This document presents a production-ready architecture for building code completion systems that power AI-assisted development tools, focusing on scalable neural models, efficient inference pipelines, and seamless integration with IDEs.

1. Problem Statement & ML Task Definition

Clarifying questions

What is the target programming language ecosystem? Are we building for Python, JavaScript, Java, or multi-language support?
What is the expected QPS and scale? What's the target user base (e.g., 1M daily active users) and concurrent completion requests?
What is the latency budget for code completion? Is sub-100ms preferred or can we accept up to 500ms for better accuracy?
Is this an online or batch system? Should completions be generated in real-time as the user types, or can we batch and precompute?
What is the success definition? Is it measured by completion accuracy, user adoption rate, or code quality improvements?
How much training data is available? Is there access to large-scale open-source code repositories or proprietary codebases?
What is the expected model size and inference cost? Will this run on edge devices or cloud-only?

Business use case and success criteria

The system aims to provide intelligent code completion suggestions to developers in real-time, improving coding efficiency and reducing boilerplate. The primary business use case is to increase developer productivity by reducing context switching, minimizing syntax errors, and accelerating feature development cycles.

Success criteria include:

Completion Accuracy: Percentage of completions that match the developer's intent (measured via human evaluation)
User Adoption Rate: % of users who continue using the system after initial setup
Productivity Gain: Average time saved per coding session (measured via user surveys and usage analytics)
Code Quality: Reduction in bugs or code review cycles (measured via internal code quality metrics)

ML problem formulation

This is a sequence-to-sequence generation task with the following characteristics:

Task Type: Text generation / code completion (token-level autoregressive prediction)
Input: Context window of code (e.g., 500 tokens including surrounding code, comments, and identifiers)
Output: Next token or sequence of tokens that complete the code snippet
Justification:
- Code completion is fundamentally about predicting the next logical token in a sequence
- The task requires understanding context, syntax, and semantic structure of code
- Unlike classification or ranking, this is a generative task where the model must produce novel code elements

ML evaluation metrics

Offline Evaluation Metrics (Model Development):

BLEU Score: Measures n-gram overlap between generated and reference completions
ROUGE Score: Evaluates recall of generated text against reference (useful for longer completions)
Perplexity: Log-likelihood of the model on held-out test data (lower is better)
Edit Distance: Measures structural similarity between generated and reference code

Online Evaluation Metrics (Production):

Completion Acceptance Rate: % of completions that are accepted by developers (measured via user interaction logs)
First-Suggestion Accuracy: % of completions that match the developer's intent on first suggestion
Time to Completion: Average time from request to acceptance (lower is better)
Engagement Metrics: Number of completions per session, usage frequency

Scale and constraints

Back-of-envelope estimates:

Training data size: 500GB of code (e.g., GitHub repositories, StackOverflow posts, open-source projects)
Number of features: ~10,000 unique tokens (vocabulary size) + context window features (~500 tokens)
Embedding matrix memory footprint: 10,000 tokens × 512 embedding dimension = ~2MB (per model)
Serving QPS math:
- Target: 10,000 QPS (10k concurrent users with 1 completion request per second)
- Latency budget: 100ms p99 (50ms p50)
- Model size: ~1GB (for transformer model)
- Estimated inference time per request: 50ms (with optimizations)

Non-ML baseline strategy

The system should include:

Rule-based fallback: For common patterns like for, if, while loops, and standard library functions (e.g., import, def)
Template-based completions: Predefined templates for common code patterns (e.g., try-except, class definitions)
Cold start handling: For new users or languages, fall back to generic completions or static code analysis
Static analysis integration: Use AST parsing and syntax highlighting for basic context-aware suggestions

# Rule-based fallback example for common patterns
def get_fallback_completion(context_tokens, current_token):
    if current_token == "for":
        return "for i in range(10):"
    elif current_token == "if":
        return "if condition:"
    elif current_token == "import":
        return "import "
    else:
        return None

2. Data Strategy

Training Data Sources

Public code repositories: GitHub, GitLab public datasets (e.g., CodeSearchNet, BigCode)
Stack Overflow code snippets: Structured Q&A with code examples
IDE usage logs: Real-time user interactions (with proper anonymization)
Code documentation: API references, tutorials, and examples from official sources
Third-party datasets: HuggingFace code datasets, CodeSearchNet, CodeT5 datasets

Labeling Pipeline Design

Human annotation: For high-quality training examples, particularly for rare code patterns or complex completions
Weak supervision: Using heuristics (e.g., context window matching, syntax validation) to generate pseudo-labels
Programmatic labeling: Leveraging AST parsing and semantic analysis to identify valid completions within code context

Class Imbalance Handling Strategy

The distribution of code completion patterns is highly skewed, with common patterns (e.g., for i in range(n):) vastly outnumbering rare ones. Strategy:

Oversampling: For underrepresented code patterns (e.g., specific library usage)
Weighted loss functions: Apply higher weights to rare completion types during training
Stratified sampling: Ensure balanced representation in training batches

Negative Sampling Approach

Hard negative mining: Sample completions that are syntactically valid but semantically incorrect or contextually inappropriate
Contextual negatives: Generate completions that match syntax but differ in function (e.g., wrong variable names, incorrect logic flow)
Random negatives: Sample from a pool of unrelated code snippets to ensure model learns to distinguish relevant vs irrelevant completions

Data Freshness and Update Frequency Requirements

Real-time updates: For popular libraries or frequently updated code patterns (hourly)
Batch updates: For stable code snippets and documentation (daily or weekly)
Freshness window: Keep training data within 30 days of the latest code trends to avoid obsolescence

Privacy Compliance

Anonymization: Remove all personal identifiers and sensitive code patterns (e.g., API keys, passwords)
Data retention policy: Retain training data for 6 months to 1 year before deletion or archival
GDPR/CCPA compliance: Ensure all data sources are properly licensed and user consent is obtained where required

Data Quality Validation and Monitoring

Syntax validation: Use language-specific parsers to validate training examples before ingestion
Semantic correctness: Validate that completions make logical sense in context (e.g., variable usage consistency)
Deduplication: Remove duplicate code snippets to prevent overfitting
Drift detection: Monitor for distribution shifts in training data (e.g., sudden increase in specific library usage)

import hashlib
from typing import List, Dict

class DataValidator:
    def __init__(self):
        self.seen_hashes = set()
    
    def validate_code_snippet(self, code: str) -> bool:
        # Basic syntax check (simplified example)
        try:
            compile(code, '<string>', 'exec')
            return True
        except SyntaxError:
            return False
    
    def deduplicate(self, code_snippets: List[str]) -> List[str]:
        unique_snippets = []
        for snippet in code_snippets:
            # Create hash to detect duplicates
            snippet_hash = hashlib.md5(snippet.encode()).hexdigest()
            if snippet_hash not in self.seen_hashes:
                self.seen_hashes.add(snippet_hash)
                unique_snippets.append(snippet)
        return unique_snippets

3. Feature Engineering

Feature Taxonomy

Features for code completion system can be categorized as:

User Features: Programming language preferences, IDE usage patterns, user activity frequency, coding experience level (e.g., years of experience, skill tags)
Item Features: Programming language, file type, code snippet length, function/class name, method signature, code structure complexity
Context Features: Code surrounding the completion point (n-lines before/after), indentation level, current file content, project structure, recent code changes
Interaction Features: User acceptance rate of suggestions, time spent on code completion, retry patterns, code modification after suggestion

Embedding Design

Entity Embeddings:

User ID embeddings (dimension: 128)
Language/IDE type embeddings (dimension: 64)
File extension embeddings (dimension: 32)

Text Embeddings:

Code tokens (BPE encoded) - dimension: 512
Function/class names - dimension: 256
Comments and docstrings - dimension: 256

Graph Embeddings:

Call graph embeddings for function dependencies - dimension: 128
AST-based structural embeddings - dimension: 256

Real-time vs Batch Feature Computation

Feature Type	Computation Mode	Reasoning
User activity	Real-time	For dynamic user preference adaptation
File context	Batch	Can be pre-computed from file cache
Code tokens	Real-time	Context-sensitive, requires live code state
Language patterns	Batch	Stable user language preferences

Feature Store Design

Online Store:

Low-latency Redis-based store for real-time features
TTL: 5 minutes for active session features

Offline Store:

Parquet-based feature store for training and batch processing
Partitioned by date for efficient querying

Feature Freshness and TTL

Feature	Source	Update Frequency	Store Type
User activity	Session logs	Every 30s	Online
Recent code	IDE events	Every 5s	Online
Language usage	User profile	Daily	Offline
File structure	Project cache	Hourly	Offline

Cross Features and Feature Interactions

Cross features include:

User-language combinations (user_id × language)
Context-length × token type interactions
Code complexity × user experience level

Feature interaction strategy:

Use cross-product embeddings for high-value combinations
Apply attention mechanisms to capture context-aware interactions

Feature Importance and Selection Methodology

Permutation importance: Evaluate feature impact on model loss
SHAP values: For local explainability and feature ranking
Variance-based selection: Remove low-variance features to reduce noise

import torch
import torch.nn as nn
import torch.nn.functional as F

class FeatureExtractor(nn.Module):
    def __init__(self, vocab_size=10000, embedding_dim=512):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, embedding_dim)
        self.positional_encoding = nn.Parameter(torch.randn(1024, embedding_dim))
        
    def forward(self, code_tokens, context_mask):
        # Embed tokens and add positional encoding
        embedded = self.token_embedding(code_tokens) + self.positional_encoding[:code_tokens.size(1)]
        
        # Apply attention mask for context window
        masked_embedded = embedded * context_mask.unsqueeze(-1)
        
        # Return sequence representation (e.g., mean pooling)
        return masked_embedded.mean(dim=1)

class CodeCompletionFeatureEngine:
    def __init__(self):
        self.feature_extractor = FeatureExtractor()
        
    def extract_context_features(self, code_sequence, context_window=100):
        """Extract features from code context"""
        # Tokenize and embed
        tokenized = self.tokenize_code(code_sequence)
        embedded = self.feature_extractor(tokenized)
        
        # Extract structural features (e.g., indentation, nesting level)
        structural_features = self.extract_structural_features(code_sequence)
        
        # Combine features
        combined_features = torch.cat([embedded, structural_features], dim=-1)
        
        return combined_features
    
    def tokenize_code(self, code_str):
        # Simplified tokenization - in practice would use BPE or similar
        tokens = [hash(token) % 10000 for token in code_str.split()]
        return torch.tensor(tokens[:100])  # Truncate to max length
    
    def extract_structural_features(self, code_str):
        # Extract basic structural features
        lines = code_str.split('\n')
        avg_line_length = sum(len(line) for line in lines) / len(lines) if lines else 0
        nesting_level = self.calculate_nesting_level(code_str)
        
        return torch.tensor([[avg_line_length, nesting_level]], dtype=torch.float32)
    
    def calculate_nesting_level(self, code_str):
        # Simple nesting level calculation
        return code_str.count('{') - code_str.count('}')

4. Model Architecture

Model Selection with Justification

For a code completion system, we'll use a Transformer-based language model with the following rationale:

Sequence-to-sequence modeling: Code completion requires understanding context and generating relevant code sequences, which transformers excel at
Long-range dependencies: Code has complex structural dependencies that transformers can capture better than RNNs
Parallel training: Transformers support efficient distributed training on GPUs/TPUs
Pre-trained foundation models: Models like CodeGPT, CodeT5, or GitHub Copilot's internal model can be fine-tuned for code completion

We'll implement a two-tower architecture:

Context encoder: Processes the code context (before cursor)
Completion decoder: Generates code tokens conditioned on context

Trade-offs compared to alternatives:

Logistic Regression: Cannot model sequential dependencies; insufficient for code generation
GBM: Not suitable for sequence modeling; poor generalization to new code patterns
DNNs: Limited context understanding compared to Transformers
LSTM/GRU: Struggles with very long sequences and parallelization

Architecture Diagram

graph TD
    A[Code Context] --> B[Context Encoder (Transformer)]
    C[Code Completion] --> D[Completion Decoder (Transformer)]
    B --> E[Shared Representation]
    E --> D
    D --> F[Token Generation]
    F --> G[Code Output]
    
    style A fill:#e1f5fe
    style B fill:#b3e5fc
    style C fill:#e1f5fe
    style D fill:#b3e5fc
    style E fill:#4fc3f7
    style F fill:#81d4fa
    style G fill:#4fc3f7

Two-Tower / Multi-Task / Ensemble Design

Two-Tower Architecture:

Context Tower: Encodes the preceding code context using a transformer encoder
Completion Tower: Decodes the completion tokens using a transformer decoder

Multi-task Learning Components:

Primary task: Next token prediction (code completion)
Auxiliary tasks:
- Code search (retrieval of similar code snippets)
- Code summarization (to improve context understanding)
- Bug detection (to filter out incorrect completions)

Pre-trained Model Selection and Fine-tuning Strategy

We'll use a CodeGPT-style model pre-trained on GitHub code repositories, fine-tuned for completion tasks:

Base Model: CodeGPT (12-layer transformer with 768 hidden units)
Fine-tuning approach:
- Full fine-tuning for the completion decoder
- Adapter layers for context encoder (to reduce compute cost)
- Scheduled sampling during training to improve generation quality

Model Complexity vs Latency Trade-off Analysis

Model	Parameters	Inference Time (ms)	Memory Usage (GB)	Notes
CodeGPT (small)	12M	50ms	0.8	Fast, good for mobile
CodeGPT (medium)	120M	150ms	2.5	Balanced performance
CodeGPT (large)	1B	400ms	8.0	High accuracy, latency-sensitive
CodeT5 (small)	24M	60ms	1.0	Good compression, moderate quality

Offline Evaluation Baseline Framing

Baseline: Rule-based completion (e.g., context-aware template matching) with 65% accuracy on token-level prediction
Proposed Model: Transformer-based model achieving 85%+ token accuracy on held-out test set

Core Model Code Snippet

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Config

class CodeCompletionModel(nn.Module):
    def __init__(self, model_name="gpt2", max_context_length=1024):
        super().__init__()
        self.model = GPT2LMHeadModel.from_pretrained(model_name)
        self.max_context_length = max_context_length
        
    def forward(self, input_ids, labels=None):
        # input_ids shape: (batch_size, sequence_length)
        outputs = self.model(input_ids=input_ids, labels=labels)
        return outputs
    
    def generate_completion(self, context_tokens, max_new_tokens=50):
        # Generate code completion given context
        with torch.no_grad():
            generated = self.model.generate(
                input_ids=context_tokens,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                temperature=0.7,
                top_k=50,
                top_p=0.95,
                pad_token_id=self.model.config.pad_token_id or 50256
            )
        return generated
    
    def compute_loss(self, input_ids, labels):
        # Custom loss computation with token masking
        outputs = self.model(input_ids=input_ids, labels=labels)
        return outputs.loss
    
    def score_completion(self, context_tokens, completion_tokens):
        # Score a specific completion given context
        full_input = torch.cat([context_tokens, completion_tokens], dim=1)
        # Mask the context part to only score the completion
        labels = full_input.clone()
        labels[:, :context_tokens.size(1)] = -100  # Ignore context tokens
        
        outputs = self.model(input_ids=full_input, labels=labels)
        return outputs.loss

# Example usage:
# model = CodeCompletionModel()
# context = torch.tensor([[101, 202, 303]])  # Example tokenized context
# completion = torch.tensor([[404, 505, 606]])  # Example completion tokens
# loss = model.compute_loss(context, completion)

This implementation provides:

Modular design with clear separation of training and inference paths
Scalable architecture supporting various model sizes
Production-ready with proper loss computation and generation logic
Extensible for multi-task learning components
Efficient inference with proper token handling and generation parameters

5. Serving & Inference Architecture

System Diagram

graph TD
    A[Code Editor] --> B[API Gateway]
    B --> C[Request Router]
    C --> D[Embedding Cache]
    C --> E[Retrieval Service]
    E --> F[Ranking Model]
    F --> G[Scoring Service]
    G --> H[Response Formatter]
    H --> I[Code Completion Response]

    subgraph "Online Serving Layer"
        D
        E
        F
        G
        H
    end

    subgraph "Feature Store"
        J[Online Feature Store]
        K[Offline Feature Store]
    end

    subgraph "Model Serving"
        L[TorchServe/ONNX Runtime]
        M[Model Registry]
    end

Online vs Batch Serving Decision

Online serving is required for real-time code completion where users expect sub-500ms response times. Batch serving would introduce unacceptable latency for interactive development workflows.

Two-Stage Pipeline

Candidate Retrieval: Retrieve top-k code snippets from embedding database using approximate nearest neighbor search (ANNS)
Scoring & Re-ranking: Score and re-rank candidates using transformer-based ranking model

Serving Infrastructure

Primary choice: TorchServe with ONNX Runtime

Supports both PyTorch and ONNX models
Built-in model versioning and A/B testing capabilities
Good integration with Kubernetes for scaling

Alternative consideration: Triton Inference Server

Better for multi-model serving with complex routing
More mature support for heterogeneous hardware (CPU/GPU)

Latency Budget Breakdown

Stage	p50 Latency	p99 Latency	Budget
Request Processing	5ms	10ms	15ms
Embedding Lookup	20ms	35ms	40ms
Retrieval	15ms	25ms	30ms
Ranking	30ms	50ms	60ms
Response Formatting	5ms	10ms	15ms
Total	80ms	130ms	160ms

Model Compression Strategy

Quantization:

INT8 quantization reduces model size by 4x with <2% accuracy loss for code models
Applied to embedding layers and attention weights

Pruning:

Structured pruning of attention heads (30% pruning) reduces inference time by 25%
Maintains model accuracy within 1% of full model

Distillation:

Student model trained to mimic teacher model with 20% smaller parameter count
Maintains >95% of teacher model's accuracy

Caching Strategy

Prediction Cache:

LRU cache for recent completions (10k entries)
Evicts based on usage frequency and recency

Embedding Cache:

Pre-computed embeddings for popular code snippets (1M entries)
Cached with TTL of 24 hours

KV Cache for Transformers:

For LLM-based completion models
Stores attention key/value matrices to avoid recomputation

Auto-scaling and Capacity Planning

Horizontal Scaling:

Kubernetes pods scaled based on QPS (target 500 QPS per pod)
Auto-scaler monitors p99 latency and adjusts replica count

Capacity Planning:

Peak QPS estimate: 10,000 requests/sec
Expected concurrent users: 5,000 (assuming 2 concurrent requests per user)
Target p99 latency: <150ms for 99.9% of requests

# Example inference pipeline code snippet
import torch
from transformers import AutoTokenizer, AutoModel

class CodeCompletionServing:
    def __init__(self, model_path, tokenizer_path):
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        self.model = AutoModel.from_pretrained(model_path)
        self.model.eval()
        
    def predict(self, input_tokens, max_length=200):
        with torch.no_grad():
            # Tokenize input
            inputs = self.tokenizer(
                input_tokens, 
                return_tensors="pt", 
                padding=True, 
                truncation=True
            )
            
            # Forward pass
            outputs = self.model(**inputs)
            logits = outputs.logits
            
            # Generate completion tokens
            generated_tokens = torch.argmax(logits, dim=-1)
            
            # Decode and return completion
            completion = self.tokenizer.decode(
                generated_tokens[0], 
                skip_special_tokens=True
            )
            
            return completion

# Example usage for real-time serving
serving_instance = CodeCompletionServing(
    model_path="model/checkpoint",
    tokenizer_path="tokenizer"
)

# Simulate real-time completion request
completion = serving_instance.predict(
    input_tokens="def fibonacci(n):",
    max_length=100
)

6. Evaluation Framework

Offline Evaluation Methodology

For code completion systems, we use temporal cross-validation to ensure our model generalizes across time and code evolution. We split training data chronologically, using 80% for training, 10% for validation, and 10% for testing. This approach prevents data leakage and simulates real-world deployment conditions where models encounter newer code patterns.

We also implement holdout validation with a 30-day window to capture seasonal code pattern variations and ensure robustness against temporal drift.

Slice-Based Evaluation

Critical slices for evaluation include:

User segments: New vs. experienced developers, different programming languages
Code context: Function size, nesting depth, language-specific patterns
Platform: Mobile vs. desktop, IDE type (VSCode, IntelliJ, etc.)
Time of day: Peak vs. off-peak usage patterns

A/B Test Design

Randomization unit: User-ID level randomization to prevent contamination between users.

Sample size calculation:

Power: 80%
Significance level: 5%
Minimum detectable effect: 15% improvement in completion accuracy
Expected baseline: 65% completion accuracy
Estimated sample size: ~1,200 users per variant

Test duration: 30 days to capture seasonal usage patterns and allow for statistical significance.

Primary metric: Completion accuracy (percentage of completions that match the expected code structure and syntax).

Secondary metrics:

Mean time to completion (lower is better)
User engagement with suggestions (click-through rate on top 1 suggestion)
Code quality improvement (measured via code review feedback)

Guardrail Metrics and Blocking Conditions

Guardrail thresholds:

Completion accuracy below 50% → immediate model rollback
Mean completion time increases by >30% → flag for investigation
User engagement drops below 25% of baseline → halt rollout

Blocking conditions:

More than 10% of completions generate invalid syntax
5% of completions introduce security vulnerabilities (if applicable)
Completion latency exceeds 200ms p99

Shadow Mode Rollout Plan

Phase 1 (Week 1): Shadow mode with 5% of users, logging all predictions but not serving to production
Phase 2 (Week 2): Shadow mode with 25% of users, enabling real-time logging and drift detection
Phase 3 (Week 3): Shadow mode with 50% of users, begin A/B testing with 25% of users on new model
Phase 4 (Week 4): Full rollout to 100% of users with A/B testing in progress

Long-Term Metric Measurement

Delayed conversion metrics:

User retention rate 30 days post-completion usage
Code quality improvement over time (measured via code review scores)
Developer productivity gains (measured via task completion time)

Engagement decay monitoring:

Daily active users using code completion features
Weekly/monthly engagement patterns to detect feature fatigue

# Example evaluation function for code completion accuracy
def evaluate_completion_accuracy(predictions, ground_truths):
    """
    Evaluate code completion accuracy with precision and recall metrics.
    
    Args:
        predictions: List of predicted code completions
        ground_truths: List of expected code completions
    
    Returns:
        dict: Accuracy, precision, recall metrics
    """
    correct_completions = 0
    total_predictions = len(predictions)
    
    for pred, truth in zip(predictions, ground_truths):
        # Simple exact match evaluation (could be extended with fuzzy matching)
        if pred.strip() == truth.strip():
            correct_completions += 1
    
    accuracy = correct_completions / total_predictions if total_predictions > 0 else 0
    return {
        "accuracy": accuracy,
        "total_completions": total_predictions,
        "correct_completions": correct_completions
    }

# Example A/B test evaluation logic
def ab_test_evaluation(user_ids, model_variants, completion_logs):
    """
    Evaluate A/B test results for code completion system.
    
    Args:
        user_ids: List of user identifiers
        model_variants: Dictionary mapping user_id to model variant (A/B)
        completion_logs: List of completion events with timestamps
    
    Returns:
        dict: A/B test results and statistical significance
    """
    variant_results = {"A": {"completions": 0, "accurate": 0}, 
                       "B": {"completions": 0, "accurate": 0}}
    
    for user_id, log in completion_logs.items():
        variant = model_variants.get(user_id, "A")
        if log["completion_correct"]:
            variant_results[variant]["accurate"] += 1
        variant_results[variant]["completions"] += 1
    
    # Calculate confidence intervals and p-values for statistical significance
    return variant_results

7. Training Pipeline

Training Data Pipeline

The training pipeline processes raw code repositories and generates training examples for code completion. The data flow includes:

Data Ingestion: Collect code from GitHub repositories, StackOverflow snippets, and internal codebases
Preprocessing: Tokenize code using a language-specific tokenizer (e.g., Python's tokenize module or Tree-sitter)
Context Windowing: Create training samples by sliding windows over code sequences (e.g., 512 tokens context window)
Data Shuffling: Shuffle samples across multiple shards for distributed training
Batching: Group samples into batches of 32-64 sequences for efficient training

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer

class CodeCompletionDataset(Dataset):
    def __init__(self, file_paths, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = []
        
        for file_path in file_paths:
            with open(file_path, 'r') as f:
                content = f.read()
                # Split into chunks with overlap
                chunks = self._split_into_chunks(content, max_length)
                for chunk in chunks:
                    self.samples.append(chunk)
    
    def _split_into_chunks(self, content, max_length):
        tokens = self.tokenizer.encode(content, add_special_tokens=False)
        chunks = []
        for i in range(0, len(tokens), max_length//2):
            chunk_tokens = tokens[i:i+max_length]
            chunks.append(self.tokenizer.decode(chunk_tokens, skip_special_tokens=True))
        return chunks
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        encoding = self.tokenizer(
            sample,
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten()
        }

# Data loading example
tokenizer = AutoTokenizer.from_pretrained('microsoft/CodeT5-base')
dataset = CodeCompletionDataset(['repo1.py', 'repo2.py'], tokenizer)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

Distributed Training Strategy

For large-scale training, we use data parallelism with model parallelism for the transformer layers:

Data Parallelism: Split batches across multiple GPUs (e.g., 8x A100s)
Model Parallelism: Partition transformer layers across GPUs when model doesn't fit in memory
Pipeline Parallelism: For very large models, pipeline the forward pass across multiple stages

Training Infrastructure

GPU Cluster: 8x A100 80GB GPUs for training (2-3 days per epoch)
TPU Cluster: Alternative option using 16x TPU v4s for faster training
Training Time: ~2-3 days per epoch on 8x A100s
Memory Usage: ~60GB per GPU (with gradient accumulation)

Hyperparameter Tuning Approach

We use Bayesian optimization for hyperparameter tuning:

Key Parameters:
- Learning rate (1e-4 to 1e-3)
- Batch size (32, 64, 128)
- Warmup steps (1000, 5000, 10000)
- Dropout rate (0.1, 0.2, 0.3)
- Weight decay (0.01, 0.1)

from ray import tune
from ray.tune.schedulers import ASHAScheduler

def train_function(config):
    model = CodeT5Model.from_pretrained('microsoft/CodeT5-base')
    
    # Training loop with config parameters
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config["learning_rate"],
        weight_decay=config["weight_decay"]
    )
    
    # Training code here...
    return {"accuracy": accuracy}

# Bayesian optimization setup
search_space = {
    "learning_rate": tune.loguniform(1e-4, 1e-3),
    "batch_size": tune.choice([32, 64, 128]),
    "weight_decay": tune.uniform(0.01, 0.1),
    "dropout": tune.uniform(0.1, 0.3)
}

tuner = tune.Tuner(
    train_function,
    param_space=search_space,
    tune_config=tune.TuneConfig(
        scheduler=ASHAScheduler(),
        num_samples=20
    )
)

Experiment Tracking and Reproducibility

MLflow: Track hyperparameters, metrics, model artifacts, and training logs
Version Control: Git integration to track model checkpoints and training configs
Model Registry: Centralized registry for model versioning and deployment tracking

Model Versioning and Registry Strategy

Versioning: Semantic versioning (v1.0.0, v1.1.0) with git commit hash
Registry: MLflow model registry with staging and production environments
Metadata: Training timestamp, dataset version, hardware used, performance metrics

Continual / Online Learning Design

Online Learning: Periodic retraining with new code samples (daily/weekly)
Incremental Training: Fine-tune on recent GitHub commits and popular code patterns
Feedback Loop: Incorporate user completion success rates as training signals
Data Freshness: Retrain every 7 days or when drift exceeds threshold (PSI > 0.1)

8. Monitoring & Observability

Data Drift Detection

Data drift monitoring is critical for maintaining model performance in code completion systems where programming languages, frameworks, and coding patterns evolve rapidly. We implement drift detection using:

PSI (Population Stability Index) for categorical features like programming language, framework, and code structure patterns
KL Divergence for continuous features such as token frequencies, code complexity metrics, and embedding statistics
Statistical tests (Kolmogorov-Smirnov test) for feature distribution shifts

import numpy as np
from scipy import stats

def calculate_psi(expected, actual, bins=10):
    """Calculate Population Stability Index"""
    # Normalize distributions
    expected = np.histogram(expected, bins=bins)[0]
    actual = np.histogram(actual, bins=bins)[0]
    
    # Avoid zero division
    expected = np.clip(expected, 1e-6, None)
    actual = np.clip(actual, 1e-6, None)
    
    psi = np.sum((actual - expected) * np.log(actual / expected))
    return psi

def detect_drift_features(current_data, reference_data):
    drift_metrics = {}
    for feature in current_data.columns:
        if feature in reference_data.columns:
            psi_value = calculate_psi(reference_data[feature], current_data[feature])
            drift_metrics[feature] = {
                'psi': psi_value,
                'threshold_exceeded': psi_value > 0.25
            }
    return drift_metrics

Prediction Drift and Label Drift Monitoring

Prediction drift monitors shifts in model outputs that may indicate performance degradation or concept drift:

Output distribution shift: Track distribution of predicted tokens over time using KL divergence
Top-K prediction stability: Monitor consistency in top-N token predictions for same inputs
Label drift: Track if model's predicted completions differ significantly from human-written code

Model Performance Degradation Alerting

We implement multi-tiered alerting for model performance degradation:

from datetime import datetime, timedelta
import pandas as pd

class ModelPerformanceMonitor:
    def __init__(self):
        self.performance_history = []
        self.alert_thresholds = {
            'auc_roc': 0.02,  # 2% drop triggers alert
            'ndcg': 0.05,     # 5% drop triggers alert
            'latency_p99': 0.1  # 10% latency increase triggers alert
        }
    
    def check_performance_degradation(self, current_metrics, window_hours=24):
        """Check if performance has degraded significantly"""
        recent_history = self.performance_history[-(window_hours * 2):]  # Assuming 30-min intervals
        
        if len(recent_history) < 2:
            return False
            
        baseline_metrics = recent_history[-1]
        
        degradation_detected = False
        for metric, threshold in self.alert_thresholds.items():
            if metric in current_metrics and metric in baseline_metrics:
                diff = baseline_metrics[metric] - current_metrics[metric]
                if diff > threshold:
                    degradation_detected = True
                    self.trigger_alert(f"Performance degradation detected in {metric}: {diff:.3f}")
        
        return degradation_detected
    
    def trigger_alert(self, message):
        """Log alert and notify monitoring system"""
        print(f"ALERT: {message}")
        # In production: send to Slack, PagerDuty, or email

Retraining Trigger Strategy

We implement a hybrid retraining trigger system:

Scheduled retraining: Weekly model retraining with new data
Drift-triggered: Retrain when PSI > 0.25 or NDCG drops > 5%
Performance-triggered: Retrain when QPS > 10k and latency > 500ms for 24 hours

Logging Design

Comprehensive logging ensures observability and enables debugging:

import json
import time
from typing import Dict, Any

class PredictionLogger:
    def __init__(self):
        self.log_file = "prediction_logs.jsonl"
    
    def log_prediction(self, request_id: str, input_features: Dict, 
                      predictions: Dict, model_output: str, 
                      latency_ms: float, timestamp: float):
        log_entry = {
            "request_id": request_id,
            "input_features": input_features,
            "predictions": predictions,
            "model_output": model_output,
            "latency_ms": latency_ms,
            "timestamp": timestamp,
            "user_id": input_features.get("user_id", "unknown"),
            "language": input_features.get("language", "unknown")
        }
        
        with open(self.log_file, "a") as f:
            f.write(json.dumps(log_entry) + "\n")
    
    def log_feature_update(self, feature_name: str, value: Any, 
                          timestamp: float, source: str):
        log_entry = {
            "feature_name": feature_name,
            "value": value,
            "timestamp": timestamp,
            "source": source
        }
        
        with open("feature_logs.jsonl", "a") as f:
            f.write(json.dumps(log_entry) + "\n")

Feedback Loop for Continuous Improvement

We implement a feedback collection system:

User engagement signals: Track code completion usage, acceptance rates, and edit frequency
Human-in-the-loop: Collect human corrections to improve model training data
A/B testing feedback: Use A/B test results to guide model improvements

class FeedbackCollector:
    def __init__(self):
        self.feedback_queue = []
    
    def collect_feedback(self, user_id: str, completion_id: str,
                        human_edits: list, engagement_score: float):
        feedback = {
            "user_id": user_id,
            "completion_id": completion_id,
            "human_edits": human_edits,
            "engagement_score": engagement_score,
            "timestamp": time.time()
        }
        
        # In production, queue for batch processing or real-time ingestion
        self.feedback_queue.append(feedback)
        
        # Trigger model retraining if feedback volume threshold met
        if len(self.feedback_queue) > 1000:
            self.trigger_model_update()
    
    def trigger_model_update(self):
        """Process feedback and update model training data"""
        # Implementation depends on feedback ingestion pipeline
        pass

Key Metrics to Monitor

Offline metrics: AUC-ROC, NDCG@K, BLEU score (daily)
Online metrics: CTR, completion acceptance rate, latency (real-time)
User engagement: Code reuse rate, edit frequency, session duration
System health: QPS, error rates, cache hit ratio

Alerting Cadence

Critical alerts: Real-time (latency > 500ms, performance drop > 5%)
Warning alerts: Hourly (PSI drift > 0.25, NDCG drop > 2%)
Info alerts: Daily (model version updates, training completion)

9. Responsible AI & Safety

Bias Analysis

Training data exhibits significant bias toward popular programming languages (Python, JavaScript, Java) and mainstream frameworks. The system may under-represent niche languages or less-documented code patterns, potentially creating a feedback loop that reinforces existing disparities. We must audit training data for demographic representation across:

Programming language popularity distribution
Geographic distribution of contributors (US/China/EU vs. emerging markets)
Code quality and documentation levels across different contributor groups

Fairness Metrics and Constraints

We implement fairness monitoring for:

Demographic parity: Ensure equal completion accuracy across programming language preferences (Python vs JavaScript vs Go)
Equalized odds: Maintain consistent precision/recall across different user skill levels (beginner vs intermediate vs expert)
Language fairness: Prevent systematic bias toward popular languages in completion suggestions

Explainability Requirements and Approach

For enterprise customers, we implement:

Attention visualization: Show which code context the model focuses on for each completion suggestion
SHAP feature importance: Highlight which features (code context, function signatures, variable names) drive specific completions
Code snippet explanation: For complex completions, provide human-readable rationale of the suggested code structure

LLM Safety Guardrails

For code completion systems, we implement:

Output validation: Check for syntax errors in generated completions before suggesting
Security scanning: Flag potentially dangerous code patterns (SQL injection, XSS, privilege escalation)
Rate limiting: Prevent abuse of the system for automated code generation
Content filtering: Block completions that might generate harmful or inappropriate code (e.g., malware, unethical practices)
Red teaming: Regular security testing with adversarial code examples to identify edge cases

Model Card Documentation Plan

We maintain comprehensive model cards including:

Training data sources and their distribution
Performance metrics across different programming languages and code contexts
Known limitations (e.g., handling of rare syntax, performance on edge cases)
Fairness assessment results and mitigation strategies
Safety measures implemented and tested

Incident Response Plan for Model Failures or Harmful Outputs

Immediate detection: Alert on sudden drops in completion quality or increase in syntax errors
Isolation protocol: Temporarily disable model versions showing concerning behavior
Root cause analysis: Investigate whether issues stem from training data drift, model degradation, or external factors
Rollback strategy: Maintain versioned model artifacts for quick recovery to known-good versions
User feedback loop: Provide mechanism for users to report problematic completions and flag them for review

class SafetyGuardrail:
    def __init__(self):
        self.security_patterns = [
            r"exec\(",
            r"eval\(",
            r"import.*os",
            r"import.*subprocess"
        ]
    
    def validate_completion(self, code_snippet):
        """Validate generated code for security issues"""
        for pattern in self.security_patterns:
            if re.search(pattern, code_snippet):
                return False, "Security violation detected"
        return True, "Valid completion"

    def check_syntax(self, code_snippet):
        """Validate syntax of generated completion"""
        try:
            compile(code_snippet, '<string>', 'exec')
            return True, "Syntax valid"
        except SyntaxError as e:
            return False, f"Syntax error: {str(e)}"

10. Scalability & Reliability

Bottleneck Analysis per Pipeline Stage

Training Pipeline:

Data processing: 100GB/s ingestion → 50GB/s preprocessing (bottleneck in distributed shuffling)
Model training: 200M parameters → 150GB model size (GPU memory constraint)
Feature engineering: 20M features → 50GB feature matrix (memory pressure in distributed training)

Serving Pipeline:

Embedding lookup: 100K QPS → 50ms latency target (200μs per lookup)
Model inference: 1M tokens/sec → 5ms p99 latency (20ms per token)
Feature computation: 10M concurrent users → 100ms TTL for real-time features

Back-of-envelope estimates:

Training data size: 500GB (100M code snippets)
Feature count: 20M (user, item, context features)
Embedding matrix footprint: 10GB (20M × 512 dimensions)
Serving QPS math: 100K QPS * 1ms = 100K token/sec capacity

Scale-up Strategy

Horizontal Scaling:

Serving: 1000 inference nodes with auto-scaling (200K QPS capacity)
Training: 50 GPU nodes in distributed training cluster (20x faster than single node)
Feature store: 1000 feature compute nodes (parallel feature extraction)

Distributed Serving Architecture:

Model shards across 1000 nodes (each node handles 100M parameters)
Feature store distributed across 500 nodes (partitioned by user ID)
Cache layer with 100GB memory footprint per node

Fault Tolerance

Model Serving Fallbacks:

Primary model failure → fallback to previous version (10s rollback window)
Embedding cache miss → return default embedding with confidence score
Feature store unavailability → use pre-computed fallback features (500ms latency)

Stale Model Serving:

Model version monitoring with 15-minute health checks
Automatic rollback to last stable model version on failure
Graceful degradation: return top-K completions with reduced confidence

Cold Start Handling Strategy

New User Cold Start:

Default to popular completions from trending repositories
Use user's language preference for fallback templates
Fallback to collaborative filtering with similar users (if available)

New Item Cold Start:

Default to most frequent completions in popular repositories
Use item metadata (language, framework) for content-based filtering
Fallback to ensemble of multiple models with different training data

Multi-Region Serving Strategy

Geographic Distribution:

Primary region (US-West): 60% of QPS, 80% of model capacity
Secondary region (EU): 20% of QPS, 15% of model capacity
Backup region (Asia): 20% of QPS, 5% of model capacity

Cross-region Failover:

50ms failover time budget (30ms for DNS propagation + 20ms for model load)
Regional model replicas with 15-minute sync window
Traffic steering using global load balancer with health checks

Latency Budget Allocation:

Stage	p50	p99	Budget
Request routing	1ms	5ms	5ms
Feature lookup	2ms	10ms	10ms
Model inference	5ms	20ms	20ms
Response formatting	1ms	5ms	5ms
Total	9ms	40ms	40ms

Performance Targets

Serving Latency:

95th percentile: < 50ms for completion requests
99th percentile: < 100ms for completion requests
P99.9: < 200ms (for batch operations)

Throughput Targets:

Peak QPS: 200K concurrent requests
Batch processing: 10M tokens/sec for batch jobs
Model update frequency: < 2 hours for model retraining

Reliability Targets:

SLA: 99.9% availability (1.44 minutes downtime/month)
Recovery time: < 5 minutes for model failures
Data consistency: < 1% of requests with stale features

11. Open Questions, Risks & Trade-offs

Decisions requiring product / engineering / policy alignment

User intent ambiguity: Should the system prioritize "what the user is trying to write" or "what they've typed so far"? This affects training data labeling and model objective.
Code quality vs. completion speed: Should we optimize for fastest completion or highest-quality code suggestions? This impacts model architecture and serving strategy.
Cross-language support: Should we support multiple languages in one system or build language-specific models? Affects model complexity and training data requirements.
Privacy of user code: How do we handle sensitive code patterns in training data? What's the acceptable risk for model performance?

Key ML risks

Risk	Impact	Likelihood	Mitigation
Data leakage from training code	High (model learns to reproduce exact patterns)	Medium	Implement strict data filtering and use synthetic code generation
Bias toward popular code patterns	Medium	High	Implement diversity sampling in training data and fairness-aware loss functions
Model hallucination (generating invalid code)	High	Medium	Implement syntax validation and semantic checking in inference pipeline
Performance degradation from data drift	Medium	High	Implement robust drift detection and automatic retraining triggers
Cold start for new users/items	Medium	High	Design robust fallback strategies and progressive model updates

Trade-offs made with justification

Accuracy vs. latency: Prioritize low-latency completion (≤100ms p99) over marginal accuracy gains to maintain interactive UX. This justifies model size constraints and simplified architectures.
Freshness vs. cost: Accept some staleness in training data to reduce infrastructure costs and training time, while implementing online learning for recent code patterns.
Model size vs. quality: Use model compression and quantization to meet latency requirements, accepting some accuracy trade-off for faster serving.
Online vs. batch training: Use hybrid approach with online learning for recent patterns and batch training for stable model updates to balance responsiveness and stability.

Alternatives considered and rejected

Transformer-based models: Rejected due to high latency and memory requirements for real-time completion, despite superior accuracy potential.
N-gram models: Rejected due to poor generalization and inability to capture complex code structures.
Template-based systems: Rejected due to lack of adaptability and inability to handle novel code patterns.
Multi-stage retrieval + generation: Rejected for complexity and latency overhead in favor of end-to-end models for faster response.

Technical debt acknowledged

Limited code quality validation: Current system lacks comprehensive semantic analysis, which could lead to invalid completions.
Sparse feature engineering: Current feature representation may not capture complex code dependencies, requiring future investment in graph-based features.
Monolithic model architecture: Single model for all languages and contexts may not scale well, requiring future modularization.
Limited explainability: Current model is largely black-box, making it hard to understand why certain completions are suggested.
No cross-language generalization: Models trained per language may not leverage shared code patterns across languages.

12. Phased Implementation Roadmap

Phase 1: MVP (Weeks 1-8)

Objective: Build minimal viable system with rule-based fallback and basic ML model

Deliverables:

Rule-based code completion engine (syntax-aware)
Basic token-level language model (e.g., GPT-2 with 100M parameters)
Online serving infrastructure for single model
Simple feature store with user/item embeddings

Key Tasks:

Week 1-2: Set up training infrastructure, data pipeline, and model registry
Week 3-4: Train initial token-level model with 10M code samples
Week 5-6: Deploy basic model serving with TorchServe, implement simple caching
Week 7-8: Integrate rule-based fallback and basic feature store

Dependencies:

Data pipeline ready (10M+ code samples)
Training cluster provisioned
Basic serving infrastructure in place

Phase 2: Scale (Weeks 9-16)

Objective: Scale to production-grade system with feature store and distributed training

Deliverables:

Distributed training pipeline (data parallelism across 8 GPUs)
Feature store with real-time and batch features
Multi-model ensemble (2-3 models: syntax-aware, semantic, context-aware)
A/B testing infrastructure with shadow mode
Real-time feature computation pipeline

Key Tasks:

Week 9-10: Implement distributed training with PyTorch Lightning
Week 11-12: Build feature store with online/offline sync (500K QPS target)
Week 13-14: Deploy ensemble model with candidate retrieval → scoring pipeline
Week 15-16: Integrate A/B testing framework, shadow mode rollout

Dependencies:

Feature store ready with 100M+ features
Distributed training cluster (8x GPU nodes)
A/B testing infrastructure provisioned

Phase 3: Excellence (Weeks 17-24)

Objective: Advanced features, continuous learning, and performance optimization

Deliverables:

Real-time feature computation pipeline (user context, item popularity)
Continuous learning system with drift detection and retraining triggers
Advanced model architectures (Transformer-based with attention visualization)
Performance optimization (quantization, pruning, caching strategies)
Multi-region serving capability

Key Tasks:

Week 17-18: Implement continuous learning with drift detection (PSI > 0.1)
Week 19-20: Deploy model compression (4-bit quantization, 30% size reduction)
Week 21-22: Add real-time feature computation (user activity, code context)
Week 23-24: Multi-region deployment with failover strategy and latency optimization

Dependencies:

Monitoring system in place (drift detection, performance alerts)
Model compression techniques validated offline
Multi-region infrastructure provisioned

Critical Path Dependencies:

Data Pipeline - Must be ready before model training begins (Week 2)
Feature Store - Critical for serving and training (Week 12)
Serving Infrastructure - Must support QPS targets (Week 8)
A/B Testing Framework - Required for production rollout (Week 16)

Risk Mitigation:

Data Quality: Implement data validation pipeline with drift detection
Latency: Pre-compute hot features and use efficient caching strategies
Scalability: Design for horizontal scaling from day one
Model Degradation: Implement continuous monitoring and retraining triggers

Success Metrics:

QPS: 10K+ requests/sec (p99 < 200ms)
Accuracy: NDCG@10 > 0.85 (offline), CTR > 0.02 (online)
Latency: p50 < 100ms, p99 < 200ms for completion requests
User Adoption: 30%+ completion rate increase over baseline

Resource Allocation:

Engineering: 4 ML engineers, 2 backend engineers, 1 DevOps engineer
Infrastructure: 8x GPU nodes for training, 16x CPU nodes for serving
Timeline: 24 weeks total with 3-week sprints and bi-weekly reviews