AI System Design - Data → Model → Inference → Feedback

# Production AI System Architecture
# ─────────────────────────────────────────────────────────────────

# STAGE 1: DATA PIPELINE
# Raw data → cleaned, labeled, versioned dataset
class DataPipeline:
    def __init__(self):
        self.sources = ["user_events", "database", "third_party_api"]

    def collect(self) -> list[dict]:
        """Ingest raw data from all sources."""
        raw = []
        for source in self.sources:
            raw.extend(self._fetch(source))
        return raw

    def clean(self, raw: list[dict]) -> list[dict]:
        """Remove duplicates, fix nulls, normalize formats."""
        return [r for r in raw if self._is_valid(r)]

    def label(self, clean: list[dict]) -> list[dict]:
        """Human annotation OR programmatic labeling (Snorkel-style)."""
        # For LLMs: RLHF requires human preference pairs
        return [{"input": r["text"], "label": r.get("category", "unknown")} for r in clean]

    def version(self, dataset: list[dict], version: str) -> None:
        """Store with DVC or MLflow for reproducibility."""
        print(f"Dataset v{version} saved: {len(dataset)} samples")

    def _fetch(self, source: str) -> list[dict]: return []
    def _is_valid(self, r: dict) -> bool: return bool(r)


# STAGE 2: MODEL TRAINING
# Choose architecture → train → evaluate → iterate
class ModelTraining:
    """Orchestrated by MLflow / W&B / SageMaker."""

    def experiment(self, config: dict) -> float:
        """Returns validation metric for hyperparameter search."""
        import torch
        model = self._build_model(config["architecture"])
        optimizer = torch.optim.AdamW(model.parameters(), lr=config["lr"])
        # train loop → return val_loss
        return 0.92  # example val accuracy

    def track(self, run_name: str, params: dict, metrics: dict) -> None:
        """Log to Weights & Biases or MLflow."""
        print(f"Run: {run_name} | params={params} | metrics={metrics}")

    def _build_model(self, arch: str):
        import torch.nn as nn
        return nn.Linear(768, 10)  # example


# STAGE 3: MODEL SERVING (Inference)
# REST API wrapping model - FastAPI is the standard
from fastapi import FastAPI
import torch

app = FastAPI(title="AI Inference API")
model = None  # loaded at startup

@app.on_event("startup")
async def load_model():
    global model
    model = torch.load("model.pt", map_location="cpu")
    model.eval()

@app.post("/predict")
async def predict(text: str) -> dict:
    """Single inference endpoint - <50ms p99 target."""
    with torch.no_grad():
        tokens = tokenize(text)        # preprocess
        logits = model(tokens)         # forward pass
        pred = logits.argmax(dim=-1)   # decode
    return {"prediction": pred.item(), "confidence": logits.softmax(-1).max().item()}

def tokenize(text: str): return None  # placeholder


# STAGE 4: ONLINE MONITORING
# Watch for: data drift, model degradation, latency spikes
class ModelMonitor:
    def __init__(self, baseline_accuracy: float = 0.95):
        self.baseline = baseline_accuracy
        self.alert_threshold = 0.03  # alert if drops 3%

    def check_data_drift(self, current_distribution: dict, reference: dict) -> bool:
        """Jensen-Shannon divergence between feature distributions."""
        # Use Evidently AI or Arize for production
        drift_detected = abs(current_distribution.get("mean", 0) - reference.get("mean", 0)) > 0.1
        if drift_detected:
            self.alert("DATA DRIFT DETECTED - retrigger training pipeline")
        return drift_detected

    def check_accuracy(self, current_acc: float) -> None:
        if self.baseline - current_acc > self.alert_threshold:
            self.alert(f"ACCURACY DROP: {self.baseline:.2%} → {current_acc:.2%}")

    def alert(self, msg: str) -> None:
        print(f"🚨 ALERT: {msg}")
        # Send to PagerDuty / Slack

# STAGE 5: FEEDBACK LOOP
class FeedbackLoop:
    """Human feedback → new labeled data → retrain → deploy."""

    def collect_corrections(self, predictions: list[dict]) -> list[dict]:
        """Users clicking 'wrong answer' provides implicit feedback."""
        return [p for p in predictions if p.get("user_feedback") == "incorrect"]

    def schedule_retraining(self, new_data_count: int, threshold: int = 1000) -> bool:
        if new_data_count >= threshold:
            print(f"Triggering retraining: {new_data_count} new labeled samples")
            return True
        return False