Architectures, Extensions & Practical Deployment

# Simplified Mixture of Experts routing layer.
# The router (gating network) assigns each token to top-k experts.
# Only selected experts process each token, keeping compute sparse.
import torch
import torch.nn as nn

class MoELayer(nn.Module):
    def __init__(self, d_model, n_experts, top_k=2):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k

        # Each expert is an independent FFN
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_model * 4),  # Expand
                nn.GELU(),
                nn.Linear(d_model * 4, d_model)   # Project back
            ) for _ in range(n_experts)
        ])

        # Router: lightweight linear layer that scores experts
        self.router = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):
        # x shape: (batch, seq_len, d_model)
        gate_scores = self.router(x)                        # (batch, seq, n_experts)
        top_vals, top_idx = gate_scores.topk(self.top_k)    # Select top-k experts
        weights = torch.softmax(top_vals, dim=-1)            # Normalize gate weights

        # Dispatch tokens to selected experts (simplified)
        output = torch.zeros_like(x)
        for k in range(self.top_k):
            for e in range(self.n_experts):
                mask = (top_idx[..., k] == e)                # Tokens routed to expert e
                if mask.any():
                    expert_out = self.experts[e](x[mask])    # Run expert on its tokens
                    output[mask] += weights[mask, k:k+1] * expert_out
        return output

# Load-balancing auxiliary loss to prevent router collapse.
# Encourages even distribution of tokens across experts
# by penalizing both over- and under-utilization.
import torch

def load_balancing_loss(gate_probs, top_indices, n_experts):
    """
    Compute the Switch Transformer load-balancing loss.

    gate_probs: (batch*seq, n_experts) - softmax router probabilities
    top_indices: (batch*seq,) - index of the selected expert per token
    n_experts: int - total number of experts
    """
    # f_i = fraction of tokens assigned to expert i
    one_hot = torch.nn.functional.one_hot(top_indices, n_experts).float()
    f = one_hot.mean(dim=0)  # (n_experts,)

    # p_i = average gate probability for expert i across all tokens
    p = gate_probs.mean(dim=0)  # (n_experts,)

    # Loss = n_experts * sum(f_i * p_i)
    # Minimized when f and p are uniform (1/n_experts each)
    loss = n_experts * (f * p).sum()

    return loss  # Add to main loss with a small coefficient (e.g., 0.01)

# PyTorch provides AdaptiveLogSoftmaxWithLoss for efficient large-vocabulary models.
# It partitions the vocabulary into clusters by frequency,
# using smaller projections for rare tokens.
import torch
import torch.nn as nn

# Configuration
d_model = 512       # Hidden dimension
vocab_size = 100000  # Large vocabulary (e.g., multilingual)

# Define cluster cutoffs: [2000, 10000, 100000]
# - Head: tokens 0-1999 (most frequent, full-dim projection)
# - Tail 1: tokens 2000-9999 (reduced-dim projection)
# - Tail 2: tokens 10000-99999 (smallest projection)
adaptive_softmax = nn.AdaptiveLogSoftmaxWithLoss(
    in_features=d_model,
    n_classes=vocab_size,
    cutoffs=[2000, 10000],  # Frequency-based cluster boundaries
    div_value=4.0             # Dimension reduction factor for tail clusters
)

# Forward pass: computes log-probabilities and loss efficiently
hidden = torch.randn(32, d_model)   # Batch of hidden states
targets = torch.randint(0, vocab_size, (32,))  # Target token ids

output = adaptive_softmax(hidden, targets)
print(f"Loss: {output.loss.item():.4f}")

# For inference: get full log-probabilities (when needed)
log_probs = adaptive_softmax.log_prob(hidden)  # (batch, vocab_size)

# Query a knowledge graph and inject structured facts into an LLM prompt.
# This pattern grounds the model in verified entities and relations,
# reducing hallucination on factual questions.
from neo4j import GraphDatabase

class KGAugmenter:
    def __init__(self, uri, auth):
        self.driver = GraphDatabase.driver(uri, auth=auth)

    def get_entity_context(self, entity_name: str, max_hops: int = 2) -> str:
        """Retrieve structured facts about an entity from the knowledge graph."""
        query = """
        MATCH (e {name: $name})-[r*1..%d]-(related)
        RETURN e.name, type(r[0]), related.name
        LIMIT 20
        """ % max_hops

        with self.driver.session() as session:
            results = session.run(query, name=entity_name)
            # Format triples as structured context for the LLM
            facts = []
            for record in results:
                facts.append(
                    f"({record[0]}) --[{record[1]}]--> ({record[2]})"
                )
        return "\n".join(facts)

    def build_augmented_prompt(self, question: str, entity: str) -> str:
        """Combine KG facts with the user question."""
        facts = self.get_entity_context(entity)
        return f"""Known facts from our knowledge graph:
{facts}

Based on these verified facts, answer: {question}"""

# Hybrid retrieval: vector search finds relevant passages,
# then knowledge graph validates and enriches entity references.
# This gives you broad coverage + entity precision.

class HybridRetriever:
    def __init__(self, vector_store, kg_client, entity_extractor):
        self.vectors = vector_store       # e.g., Pinecone, Weaviate, pgvector
        self.kg = kg_client               # e.g., Neo4j, Amazon Neptune
        self.extractor = entity_extractor  # NER model for entity detection

    def retrieve(self, query: str, top_k: int = 5) -> dict:
        """Two-stage retrieval: vector search + KG enrichment."""

        # Stage 1: Vector retrieval for broad evidence discovery
        passages = self.vectors.search(query, top_k=top_k)

        # Stage 2: Extract entities mentioned in query and passages
        entities = self.extractor.extract(query)
        for p in passages:
            entities.extend(self.extractor.extract(p.text))
        entities = list(set(entities))  # Deduplicate

        # Stage 3: KG lookup for structured facts about those entities
        kg_facts = []
        for entity in entities:
            facts = self.kg.get_relations(entity, max_hops=2)
            kg_facts.extend(facts)

        return {
            "passages": passages,     # Semantic evidence (from vectors)
            "facts": kg_facts,         # Structured facts (from KG)
            "entities": entities       # Detected entities
        }

# A simple evaluation harness for comparing models empirically.
# Runs the same test suite against multiple providers and
# reports quality, latency, and cost metrics.
import time
from dataclasses import dataclass

@dataclass
class EvalResult:
    model: str
    task: str
    score: float        # Task-specific quality metric (0-1)
    ttft_ms: float      # Time to first token
    total_ms: float     # Total generation time
    cost_usd: float     # Estimated cost per request

class ModelComparator:
    def __init__(self, models, evaluator, test_suite):
        self.models = models         # Dict of model_name -> client
        self.evaluator = evaluator   # Scores output quality per task
        self.test_suite = test_suite # List of (task_name, prompt, expected)

    def run_comparison(self) -> list:
        """Run all test cases against all models."""
        results = []
        for model_name, client in self.models.items():
            for task, prompt, expected in self.test_suite:
                # Measure timing
                start = time.perf_counter()
                output = client.generate(prompt)
                elapsed = (time.perf_counter() - start) * 1000

                # Score quality
                score = self.evaluator.score(output, expected, task)

                results.append(EvalResult(
                    model=model_name, task=task,
                    score=score, ttft_ms=output.ttft_ms,
                    total_ms=elapsed, cost_usd=output.usage_cost
                ))
        return results

# Define a complete hyperparameter configuration for a RAG system.
# This shows that production LLM tuning goes far beyond learning rate.
from dataclasses import dataclass

@dataclass
class RAGConfig:
    """Full hyperparameter surface for a RAG-based LLM system."""

    # --- Retrieval hyperparameters ---
    chunk_size: int = 512            # Tokens per chunk (affects granularity)
    chunk_overlap: int = 64          # Overlap between chunks (context continuity)
    retrieval_top_k: int = 10        # Candidates from vector search
    reranker_top_n: int = 3          # Final candidates after reranking
    similarity_threshold: float = 0.7  # Minimum relevance score

    # --- Decoding hyperparameters ---
    temperature: float = 0.3        # Low for factual RAG tasks
    top_p: float = 0.85              # Nucleus sampling threshold
    max_output_tokens: int = 1024   # Output budget (cost + latency)

    # --- Model hyperparameters ---
    model_name: str = "claude-sonnet"   # Model selection
    context_budget: int = 8000       # Max tokens for retrieved context

    # --- Fine-tuning hyperparameters (if applicable) ---
    lora_rank: int = 16              # LoRA rank (adaptation capacity)
    lora_alpha: float = 32.0         # LoRA scaling factor
    weight_decay: float = 0.01       # Regularization

# Example: create configs for different use cases
factual_config = RAGConfig(temperature=0.1, retrieval_top_k=20, reranker_top_n=5)
creative_config = RAGConfig(temperature=0.8, top_p=0.95, retrieval_top_k=5)

# Disaggregated evaluation: break overall metrics into subgroup metrics
# to expose hidden bias that aggregate numbers mask.
from collections import defaultdict

class DisaggregatedEvaluator:
    def __init__(self, dimensions):
        """
        dimensions: list of metadata keys to disaggregate by
                    e.g., ["demographic", "topic", "language", "complexity"]
        """
        self.dimensions = dimensions
        self.results = defaultdict(lambda: defaultdict(list))

    def record(self, prediction, ground_truth, metadata: dict):
        """Record one evaluation result with its metadata."""
        is_correct = self._score(prediction, ground_truth)

        # Record overall
        self.results["_overall"]["all"].append(is_correct)

        # Record per dimension
        for dim in self.dimensions:
            if dim in metadata:
                self.results[dim][metadata[dim]].append(is_correct)

    def report(self):
        """Generate disaggregated performance report."""
        for dim, groups in self.results.items():
            print(f"\n--- {dim} ---")
            for group, scores in sorted(groups.items()):
                acc = sum(scores) / len(scores)
                print(f"  {group}: {acc:.1%} ({len(scores)} samples)")

                # Flag significant disparities
                overall_acc = sum(self.results["_overall"]["all"]) / \
                              len(self.results["_overall"]["all"])
                if abs(acc - overall_acc) > 0.1:
                    print(f"    WARNING: {abs(acc-overall_acc):.1%} gap vs overall")

# Privacy-aware logging middleware that captures metrics
# without storing sensitive content. This balances
# observability needs with data protection requirements.
import hashlib
import time
from dataclasses import dataclass

@dataclass
class PrivacySafeLog:
    """Log entry that captures operational metrics without PII."""
    request_hash: str     # One-way hash for deduplication, not content
    timestamp: float
    model: str
    input_tokens: int      # Count only, not content
    output_tokens: int
    ttft_ms: float
    total_ms: float
    safety_score: float    # Safety classifier output
    pii_detected: bool     # Whether PII was found (not what it was)
    user_tier: str          # Anonymized user group, not user ID

class PrivacyLogger:
    def __init__(self, pii_detector, log_store):
        self.pii_det = pii_detector
        self.store = log_store

    def log_request(self, prompt, response, metadata):
        """Log metrics about a request without storing content."""
        # Hash the prompt for deduplication (cannot be reversed)
        prompt_hash = hashlib.sha256(prompt.encode()).hexdigest()[:16]

        # Check for PII without logging what was found
        has_pii = self.pii_det.contains_pii(response)

        # Store only safe metadata
        entry = PrivacySafeLog(
            request_hash=prompt_hash,
            timestamp=time.time(),
            model=metadata["model"],
            input_tokens=metadata["input_tokens"],
            output_tokens=metadata["output_tokens"],
            ttft_ms=metadata["ttft_ms"],
            total_ms=metadata["total_ms"],
            safety_score=metadata["safety_score"],
            pii_detected=has_pii,
            user_tier=metadata.get("user_tier", "unknown")
        )
        self.store.write(entry)
        # Content (prompt, response) is NEVER stored

# Deployment readiness checklist that codifies the governance matrix.
# Run this before any production deployment to ensure all areas are covered.
from dataclasses import dataclass, field

@dataclass
class DeploymentCheck:
    name: str
    area: str         # privacy | quality | safety | operations
    passed: bool = False
    notes: str = ""

class DeploymentReadiness:
    def __init__(self):
        self.checks = [
            # Privacy controls
            DeploymentCheck("PII detection enabled", "privacy"),
            DeploymentCheck("Access control integrated with retrieval", "privacy"),
            DeploymentCheck("Logging policy reviewed (no content in logs)", "privacy"),
            DeploymentCheck("Data retention policy configured", "privacy"),

            # Quality controls
            DeploymentCheck("Regression test suite passes", "quality"),
            DeploymentCheck("Eval suite versioned with prompts", "quality"),
            DeploymentCheck("A/B test plan defined", "quality"),
            DeploymentCheck("Prompt pinned to model version", "quality"),

            # Safety controls
            DeploymentCheck("Input safety classifier deployed", "safety"),
            DeploymentCheck("Output moderation enabled", "safety"),
            DeploymentCheck("Human escalation path configured", "safety"),
            DeploymentCheck("Red-team results reviewed", "safety"),

            # Operations controls
            DeploymentCheck("Tracing enabled (end-to-end request IDs)", "operations"),
            DeploymentCheck("Rollback procedure tested", "operations"),
            DeploymentCheck("SLO dashboards configured", "operations"),
            DeploymentCheck("Cost attribution tags in place", "operations"),
        ]

    def report(self):
        """Print deployment readiness status."""
        for area in ["privacy", "quality", "safety", "operations"]:
            area_checks = [c for c in self.checks if c.area == area]
            passed = sum(1 for c in area_checks if c.passed)
            print(f"\n{area.upper()}: {passed}/{len(area_checks)}")
            for c in area_checks:
                status = "PASS" if c.passed else "FAIL"
                print(f"  [{status}] {c.name}")