Upload folder using huggingface_hub

Dockerfile

@@ -0,0 +1,27 @@
+FROM python:3.11-slim
+
+# Install system dependencies (needed for git-lfs and health checks)
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    curl \
+    git \
+    git-lfs \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Copy requirements and install torch CPU version
+COPY requirements.txt .
+RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy all source files
+COPY . .
+
+# Create logs directory
+RUN mkdir -p logs
+
+# Expose the default port for Hugging Face Spaces
+EXPOSE 7860
+
+# Run uvicorn server on port 7860
+CMD ["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,10 +1,19 @@
 ---
-title: Nexa Classify Api
-emoji: 🐢
-colorFrom: purple
-colorTo: blue
+title: Nexa Classify API
+emoji: 🚀
+colorFrom: indigo
+colorTo: purple
 sdk: docker
+app_port: 7860
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Nexa Classify FastAPI Backend
+
+This is the production-ready FastAPI backend for document classification, containerized for Hugging Face Spaces.
+
+## Setup & Running locally
+If you want to run this backend locally, use:
+```bash
+conda run -n docclassifier uvicorn api:app --host 0.0.0.0 --port 8000 --reload
+```

api.py

@@ -0,0 +1,456 @@
+"""
+api.py
+──────
+Production-ready REST API for document classification.
+Supports any combination of saved models via lazy loading.
+
+Usage
+─────
+    pip install fastapi uvicorn[standard] pydantic   (already in requirements.txt)
+
+    # Start the server (from project root, venv active)
+    uvicorn api:app --host 0.0.0.0 --port 8000 --reload
+
+    # Health check
+    curl http://localhost:8000/health
+
+    # Single prediction
+    curl -X POST http://localhost:8000/predict \
+      -H "Content-Type: application/json" \
+      -d '{"text": "Fed raises interest rates by 50 bps", "model_name": "roberta_base"}'
+
+    # Batch prediction
+    curl -X POST http://localhost:8000/batch_predict \
+      -H "Content-Type: application/json" \
+      -d '{"texts": ["Apple unveils M5 chip", "Ronaldo scores again"], "model_name": "roberta_base"}'
+
+    # Explore interactive docs at: http://localhost:8000/docs
+"""
+import logging
+import os
+import time
+from contextlib import asynccontextmanager
+from typing import Dict, List, Optional
+from uuid import uuid4
+
+import numpy as np
+import torch
+from fastapi import FastAPI, HTTPException, Request
+from pydantic import BaseModel, Field
+
+from config import CFG
+import database
+
+logger = logging.getLogger("api")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(levelname)s  %(message)s")
+
+
+# ── Pydantic schemas ──────────────────────────────────────────────────────────
+
+class PredictRequest(BaseModel):
+    text: str       = Field(..., min_length=1, max_length=10_000,
+                            example="Apple launches a groundbreaking AI chip.")
+    model_name: str = Field(default="roberta_base",
+                            example="roberta_base",
+                            description="Directory name in saved_models/. "
+                                        "Examples: 'roberta_base', 'lr', 'svm'.")
+
+
+class BatchPredictRequest(BaseModel):
+    texts:      List[str] = Field(..., min_length=1, max_length=256)
+    model_name: str       = Field(default="roberta_base")
+
+
+class Prediction(BaseModel):
+    text:          str
+    request_id:    str
+    label_id:      int
+    label:         str
+    probabilities: Optional[Dict[str, float]] = None
+    is_low_confidence: bool
+    latency_ms:    float
+
+
+class BatchResponse(BaseModel):
+    predictions:     List[Prediction]
+    count:           int
+    total_latency_ms: float
+
+
+class HealthResponse(BaseModel):
+    status:        str
+    loaded_models: List[str]
+    quantized:     Dict[str, bool]
+    device:        str
+    version:       str = "1.0.0"
+
+
+# ── Model registry ────────────────────────────────────────────────────────────
+
+_registry: Dict[str, Dict] = {}   # model_name → {"obj": ..., "kind": str, "quantized": bool}
+
+
+def _load_model(model_name: str):
+    """Lazy-load a model on first request, then cache it in _registry."""
+    # Normalise model name mapping (case-insensitive & support aliases)
+    name_lower = model_name.lower()
+    if name_lower in ("lr", "svm"):
+        model_name = name_lower
+    elif "distilbert" in name_lower:
+        model_name = "distilbert_base_uncased"
+    elif "roberta" in name_lower:
+        model_name = "roberta_base"
+    elif "bert" in name_lower:
+        model_name = "bert_base_uncased"
+
+    if model_name in _registry:
+        entry = _registry[model_name]
+        return entry["obj"], entry["kind"], entry["quantized"]
+
+    if model_name in ("lr", "svm"):
+        import joblib
+        path = os.path.join(CFG.models_dir, f"traditional_{model_name}.joblib")
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"No model file: {path}")
+        obj  = joblib.load(path)
+        kind = "sklearn"
+        quantized = False
+    else:
+        from transformers import AutoModelForSequenceClassification, AutoTokenizer
+        from transformer_model import _checkpoint_to_dir
+
+        path = os.path.join(CFG.models_dir, model_name)
+        if not os.path.isdir(path):
+            alt = os.path.join(CFG.models_dir, _checkpoint_to_dir(model_name))
+            if os.path.isdir(alt):
+                path = alt
+            else:
+                raise FileNotFoundError(
+                    f"No model directory: {path}\n"
+                    f"Hint: check saved_models/ for available directories."
+                )
+
+        int8_path = f"{path}_int8"
+        int8_file = os.path.join(int8_path, "model_int8.pt")
+        if os.path.exists(int8_file):
+            try:
+                torch.backends.quantized.engine = "qnnpack"
+            except Exception:
+                pass
+            try:
+                model = torch.load(int8_file, map_location="cpu", weights_only=False)
+            except TypeError:
+                model = torch.load(int8_file, map_location="cpu")
+            tokenizer = AutoTokenizer.from_pretrained(int8_path)
+            quantized = True
+        else:
+            model = AutoModelForSequenceClassification.from_pretrained(path)
+            tokenizer = AutoTokenizer.from_pretrained(path)
+            quantized = False
+
+        model.eval()
+        obj  = (model, tokenizer)
+        kind = "transformer"
+
+    _registry[model_name] = {"obj": obj, "kind": kind, "quantized": quantized}
+    q = "int8" if quantized else "fp32"
+    logger.info(f"Model cached: {model_name}  [{kind}:{q}]")
+    return obj, kind, quantized
+
+
+def _infer_single(text: str, obj, kind: str) -> Dict:
+    if kind == "transformer":
+        model, tokenizer = obj
+        enc = tokenizer(text, truncation=True,
+                        max_length=CFG.max_length, return_tensors="pt")
+        with torch.no_grad():
+            probs = torch.softmax(model(**enc).logits[0], dim=-1).numpy()
+        pred_id = int(np.argmax(probs))
+        conf = float(np.max(probs))
+        return {
+            "label_id": pred_id,
+            "label":    CFG.label_names[pred_id],
+            "probabilities": {
+                CFG.label_names[i]: round(float(p), 4)
+                for i, p in enumerate(probs)
+            },
+            "confidence": conf,
+        }
+    # sklearn
+    pred_id = int(obj.predict([text])[0])
+    result  = {"label_id": pred_id, "label": CFG.label_names[pred_id],
+               "probabilities": None, "confidence": 1.0}
+    clf = list(obj.named_steps.values())[-1]
+    if hasattr(clf, "predict_proba"):
+        probs = obj.predict_proba([text])[0]
+        result["probabilities"] = {
+            CFG.label_names[i]: round(float(p), 4) for i, p in enumerate(probs)
+        }
+        result["confidence"] = float(np.max(probs))
+    elif hasattr(clf, "decision_function"):
+        scores = obj.decision_function([text])
+        scores = np.asarray(scores, dtype=np.float64).reshape(1, -1)
+        scores = scores - np.max(scores, axis=1, keepdims=True)
+        exps = np.exp(scores)
+        probs = exps / np.sum(exps, axis=1, keepdims=True)
+        result["confidence"] = float(np.max(probs))
+    return result
+
+
+def _infer_batch(texts: List[str], obj, kind: str) -> List[Dict]:
+    if kind == "transformer":
+        model, tokenizer = obj
+        results    = []
+        batch_size = 16
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i : i + batch_size]
+            enc   = tokenizer(batch, truncation=True, max_length=CFG.max_length,
+                              padding=True, return_tensors="pt")
+            with torch.no_grad():
+                logits = model(**enc).logits
+            probs_batch = torch.softmax(logits, dim=-1).numpy()
+            for text, probs in zip(batch, probs_batch):
+                pred_id = int(np.argmax(probs))
+                conf = float(np.max(probs))
+                results.append({
+                    "label_id": pred_id,
+                    "label":    CFG.label_names[pred_id],
+                    "probabilities": {
+                        CFG.label_names[i]: round(float(p), 4)
+                        for i, p in enumerate(probs)
+                    },
+                    "confidence": conf,
+                    "text": text,
+                })
+        return results
+    # sklearn batch
+    preds = obj.predict(texts)
+    clf = list(obj.named_steps.values())[-1]
+    confidences = np.ones(len(texts), dtype=np.float64)
+    if hasattr(clf, "predict_proba"):
+        probs = obj.predict_proba(texts)
+        confidences = np.max(probs, axis=1)
+    elif hasattr(clf, "decision_function"):
+        scores = obj.decision_function(texts)
+        scores = np.asarray(scores, dtype=np.float64)
+        if scores.ndim == 1:
+            scores = np.stack([-scores, scores], axis=1)
+        scores = scores - np.max(scores, axis=1, keepdims=True)
+        exps = np.exp(scores)
+        probs = exps / np.sum(exps, axis=1, keepdims=True)
+        confidences = np.max(probs, axis=1)
+
+    results = []
+    for p, t, c in zip(preds, texts, confidences):
+        results.append(
+            {
+                "label_id": int(p),
+                "label": CFG.label_names[int(p)],
+                "probabilities": None,
+                "confidence": float(c),
+                "text": t,
+            }
+        )
+    return results
+
+
+# ── FastAPI app ───────────────────────────────────────────────────────────────
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Pre-warm the default model on server startup."""
+    try:
+        database.init_db()
+        _load_model("roberta_base")
+        logger.info("Default model (roberta_base) pre-loaded.")
+    except FileNotFoundError:
+        logger.warning("Default model not found; will load on first request.")
+    yield
+    _registry.clear()
+    logger.info("Model registry cleared.")
+
+
+app = FastAPI(
+    title="Document Classifier API",
+    description=(
+        "Multi-class news text classification over four categories: "
+        "World · Sports · Business · Sci/Tech. "
+        "Supports traditional ML and transformer models."
+    ),
+    version="1.0.0",
+    lifespan=lifespan,
+)
+
+from fastapi.middleware.cors import CORSMiddleware
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.middleware("http")
+async def add_private_network_header(request: Request, call_next):
+    response = await call_next(request)
+    if "access-control-request-private-network" in request.headers:
+        response.headers["Access-Control-Allow-Private-Network"] = "true"
+        origin = request.headers.get("origin")
+        if origin:
+            response.headers["Access-Control-Allow-Origin"] = origin
+    return response
+
+
+
+@app.get("/health", response_model=HealthResponse, tags=["Status"])
+async def health():
+    """Confirm the API is running, list loaded models, and report device."""
+    return HealthResponse(
+        status="ok",
+        loaded_models=list(_registry.keys()),
+        quantized={k: bool(v.get("quantized")) for k, v in _registry.items()},
+        device=CFG.device,
+    )
+
+
+@app.get("/labels", tags=["Status"])
+async def get_labels():
+    """Return the four classification labels."""
+    return {
+        "labels": [
+            {"id": i, "name": n} for i, n in enumerate(CFG.label_names)
+        ]
+    }
+
+
+@app.get("/models", tags=["Status"])
+async def list_available_models():
+    """List all models that exist in saved_models/ and are ready to load."""
+    available = []
+    if os.path.isdir(CFG.models_dir):
+        for name in os.listdir(CFG.models_dir):
+            path = os.path.join(CFG.models_dir, name)
+            if name.endswith("_int8"):
+                continue
+            if os.path.isdir(path) and os.path.exists(
+                os.path.join(path, "config.json")
+            ):
+                int8_file = os.path.join(f"{path}_int8", "model_int8.pt")
+                available.append(
+                    {
+                        "name": name,
+                        "type": "transformer",
+                        "quantized": bool(os.path.exists(int8_file)),
+                    }
+                )
+        for fname in os.listdir(CFG.models_dir):
+            if fname.startswith("traditional_") and fname.endswith(".joblib"):
+                short = fname.replace("traditional_", "").replace(".joblib", "")
+                available.append({"name": short, "type": "sklearn", "quantized": False})
+    return {"models": available, "count": len(available)}
+
+
+@app.post("/predict", response_model=Prediction, tags=["Inference"])
+async def predict(req: PredictRequest):
+    """Classify a single text document and return label + probabilities."""
+    t0 = time.perf_counter()
+    request_id = str(uuid4())
+    try:
+        obj, kind, _ = _load_model(req.model_name)
+    except FileNotFoundError as exc:
+        raise HTTPException(status_code=404, detail=str(exc))
+    result  = _infer_single(req.text, obj, kind)
+    latency = (time.perf_counter() - t0) * 1000
+    confidence = float(result.get("confidence", 1.0))
+    is_low = bool(confidence < float(CFG.low_confidence_threshold))
+    database.log_request(
+        request_id=request_id,
+        model_name=req.model_name,
+        input_text=req.text,
+        predicted_label=str(result["label"]),
+        predicted_label_id=int(result["label_id"]),
+        confidence=confidence,
+        latency_ms=float(latency),
+        is_batch=False,
+    )
+    return Prediction(
+        text=req.text[:200],
+        request_id=request_id,
+        is_low_confidence=is_low,
+        latency_ms=round(latency, 2),
+        label_id=result["label_id"],
+        label=result["label"],
+        probabilities=result.get("probabilities"),
+    )
+
+
+@app.post("/batch_predict", response_model=BatchResponse, tags=["Inference"])
+async def batch_predict(req: BatchPredictRequest):
+    """Classify a list of documents in one call (up to 256 texts)."""
+    t0 = time.perf_counter()
+    try:
+        obj, kind, _ = _load_model(req.model_name)
+    except FileNotFoundError as exc:
+        raise HTTPException(status_code=404, detail=str(exc))
+    raw_results = _infer_batch(req.texts, obj, kind)
+    total_ms    = (time.perf_counter() - t0) * 1000
+    per_item_ms = (total_ms / len(req.texts)) if req.texts else 0.0
+    predictions = [
+        Prediction(
+            text=r["text"][:200],
+            request_id=str(uuid4()),
+            label_id=r["label_id"],
+            label=r["label"],
+            probabilities=r.get("probabilities"),
+            is_low_confidence=bool(float(r.get("confidence", 1.0)) < float(CFG.low_confidence_threshold)),
+            latency_ms=round(per_item_ms, 2),
+        )
+        for r in raw_results
+    ]
+    for r, pred in zip(raw_results, predictions):
+        database.log_request(
+            request_id=pred.request_id,
+            model_name=req.model_name,
+            input_text=r["text"],
+            predicted_label=str(r["label"]),
+            predicted_label_id=int(r["label_id"]),
+            confidence=float(r.get("confidence", 1.0)),
+            latency_ms=float(per_item_ms),
+            is_batch=True,
+        )
+    return BatchResponse(
+        predictions=predictions,
+        count=len(predictions),
+        total_latency_ms=round(total_ms, 2),
+    )
+
+
+@app.get("/analytics/summary", tags=["Analytics"])
+async def analytics_summary(model_name: Optional[str] = None, days: int = 7):
+    return database.get_summary(model_name=model_name, days=days)
+
+
+@app.get("/analytics/history", tags=["Analytics"])
+async def analytics_history(limit: int = 50, offset: int = 0):
+    return database.get_request_history(limit=limit, offset=offset)
+
+
+@app.get("/analytics/low_confidence", tags=["Analytics"])
+async def analytics_low_confidence(reviewed: bool = False, limit: int = 50):
+    return database.get_low_confidence_flags(reviewed=reviewed, limit=limit)
+
+
+class ReviewBody(BaseModel):
+    note: Optional[str] = None
+
+
+@app.patch("/analytics/review/{request_id}", tags=["Analytics"])
+async def analytics_mark_reviewed(request_id: str, body: ReviewBody):
+    database.mark_reviewed(request_id=request_id, note=body.note)
+    return {"request_id": request_id, "reviewed": True}
+
+
+@app.post("/analytics/export_flags", tags=["Analytics"])
+async def analytics_export_flags():
+    return database.export_low_confidence_to_folder()

compare_results.py

@@ -0,0 +1,132 @@
+"""
+compare_results.py
+──────────────────
+Auto-discovers and evaluates all saved models side-by-side.
+Handles multiple transformer architectures in saved_models/.
+
+Usage
+─────
+    python compare_results.py
+"""
+import logging
+import os
+from typing import Dict, List
+
+import numpy as np
+import torch
+from sklearn.metrics import accuracy_score, f1_score
+
+from config import CFG
+from data_loader import load_test_only
+import traditional_model as tm
+import transformer_model as trm
+
+logging.basicConfig(level=logging.WARNING)
+
+
+# ── Discovery ─────────────────────────────────────────────────────────────────
+
+def _discover_transformer_models() -> List[str]:
+    """Return directory names of all saved transformer models."""
+    found = []
+    if not os.path.isdir(CFG.models_dir):
+        return found
+    for name in sorted(os.listdir(CFG.models_dir)):
+        path = os.path.join(CFG.models_dir, name)
+        if os.path.isdir(path) and os.path.exists(os.path.join(path, "config.json")):
+            found.append(name)
+    return found
+
+
+# ── Evaluation ────────────────────────────────────────────────────────────────
+
+def _eval_traditional(name: str, X_test: List[str], y_test: List[int]) -> Dict:
+    try:
+        pipeline = tm.load_model(name)
+        preds    = list(pipeline.predict(X_test))
+        return {
+            "accuracy": accuracy_score(y_test, preds),
+            "f1_macro": f1_score(y_test, preds, average="macro"),
+        }
+    except FileNotFoundError:
+        return {}
+
+
+def _eval_transformer(model_dir: str, X_test: List[str], y_test: List[int]) -> Dict:
+    path = os.path.join(CFG.models_dir, model_dir)
+    try:
+        from transformers import AutoModelForSequenceClassification, AutoTokenizer
+        model     = AutoModelForSequenceClassification.from_pretrained(path)
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        model.eval()
+
+        preds      = []
+        batch_size = 32
+
+        for i in range(0, len(X_test), batch_size):
+            batch = X_test[i : i + batch_size]
+            enc   = tokenizer(batch, truncation=True, max_length=CFG.max_length,
+                              padding=True, return_tensors="pt")
+            with torch.no_grad():
+                logits = model(**enc).logits
+            preds.extend(logits.argmax(dim=-1).tolist())
+
+        return {
+            "accuracy": accuracy_score(y_test, preds),
+            "f1_macro": f1_score(y_test, preds, average="macro"),
+        }
+    except FileNotFoundError:
+        return {}
+    except Exception as exc:
+        print(f"  [{model_dir}] Error: {exc}")
+        return {}
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+
+def main() -> None:
+    print("\n  Loading AG News test set …")
+    X_test, y_test = load_test_only()
+    print(f"  Loaded {len(X_test):,} examples.\n")
+
+    results: Dict[str, Dict] = {}
+
+    # Traditional models
+    for name in ["lr", "svm"]:
+        print(f"  Evaluating {name.upper()} …")
+        r = _eval_traditional(name, X_test, y_test)
+        if r:
+            results[name.upper()] = r
+        else:
+            print(f"    [{name.upper()}] not found — skipping.")
+
+    # All saved transformer models
+    transformer_dirs = _discover_transformer_models()
+    if not transformer_dirs:
+        print("  No transformer models found in saved_models/.")
+    for model_dir in transformer_dirs:
+        display_name = model_dir.replace("_", "-")
+        print(f"  Evaluating {display_name} …")
+        r = _eval_transformer(model_dir, X_test, y_test)
+        if r:
+            results[display_name] = r
+
+    if not results:
+        print("\n  No models found. Train at least one model first.\n")
+        return
+
+    # Print table
+    print("\n" + "═" * 58)
+    print(f"  {'Model':<22}  {'Accuracy':>10}  {'F1-Macro':>10}")
+    print("─" * 58)
+    for name, m in sorted(results.items(), key=lambda x: x[1]["accuracy"], reverse=True):
+        star = " ◀ best" if name == max(results, key=lambda k: results[k]["accuracy"]) else ""
+        print(
+            f"  {name:<22}  {m['accuracy']*100:>9.2f}%  "
+            f"{m['f1_macro']:>10.4f}{star}"
+        )
+    print("═" * 58 + "\n")
+
+
+if __name__ == "__main__":
+    main()

config.py

@@ -0,0 +1,85 @@
+"""
+config.py
+─────────
+Updated for Mac M4 / Apple Silicon MPS.
+
+Key changes vs Windows version:
+  - Auto device detection: MPS → CUDA → CPU
+  - Sample caps removed (full 120 K dataset now feasible)
+  - batch_size 16  (grad_accum_steps=2 → effective batch 32)
+  - max_length 128 (AG News headlines fit comfortably; saves ~60% VRAM vs 256)
+  - Default model upgraded to 'roberta-base'
+  - label_smoothing added for better calibration
+  - gradient_checkpointing enabled by default (MPS OOM safeguard)
+"""
+import os
+import torch
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+def _auto_device() -> str:
+    """Detect the best available compute device at import time."""
+    if torch.backends.mps.is_available() and torch.backends.mps.is_built():
+        return "mps"
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+@dataclass
+class Config:
+    # ── Dataset ──────────────────────────────────────────────────────────────
+    dataset_name: str       = "ag_news"
+    num_labels:   int       = 4
+    label_names:  List[str] = field(
+        default_factory=lambda: ["World", "Sports", "Business", "Sci/Tech"]
+    )
+
+    # Full dataset — no caps needed on M4 MPS
+    max_train_samples: Optional[int] = None   # 120,000
+    max_eval_samples:  Optional[int] = None   # ~12,000
+    max_test_samples:  Optional[int] = None   # 7,600
+
+    # ── Model ─────────────────────────────────────────────────────────────────
+    # Supported checkpoints (swap as needed):
+    #   "distilbert-base-uncased"  — 66M params  — fastest  (~45–70 min on M4)
+    #   "bert-base-uncased"        — 110M params  — balanced (~90–120 min on M4)
+    #   "roberta-base"             — 125M params  — best acc (~90–150 min on M4)
+    model_checkpoint: str = "roberta-base"
+    max_length:       int = 128    # 128 is ample for AG News; saves ~60% VRAM vs 256
+
+    # ── Training Hyper-parameters ─────────────────────────────────────────────
+    batch_size:                 int   = 16    # 16 × grad_accum_steps=2 → effective batch 32
+    num_epochs:                 int   = 3     # Safe training epochs
+    learning_rate:              float = 2e-5
+    warmup_ratio:               float = 0.06
+    weight_decay:               float = 0.01
+    grad_accum_steps:           int   = 2     # Accumulate 2 steps → effective batch 32
+    label_smoothing:            float = 0.1   # Regularisation: prevents over-confidence
+    use_gradient_checkpointing: bool  = True  # ON by default — critical MPS OOM safeguard
+
+    # ── Hardware (auto-detected) ───────────────────────────────────────────────
+    device:      str = field(default_factory=_auto_device)
+    # num_workers=0 is safest with HuggingFace datasets in torch format on Mac
+    num_workers: int = 0
+    seed:        int = 42
+    low_confidence_threshold: float = 0.60
+
+    # ── Paths ─────────────────────────────────────────────────────────────────
+    data_dir:    str = "data"
+    models_dir:  str = "saved_models"
+    outputs_dir: str = "outputs"
+    logs_dir:    str = os.path.join("outputs", "logs")
+
+    def __post_init__(self) -> None:
+        for d in [self.data_dir, self.models_dir, self.outputs_dir, self.logs_dir]:
+            os.makedirs(d, exist_ok=True)
+        device_label = (
+            "MPS — Apple Metal (M4)" if self.device == "mps" else self.device.upper()
+        )
+        print(f"[Config] Device: {device_label}  |  Model: {self.model_checkpoint}")
+
+
+# Module-level singleton — imported by all modules
+CFG = Config()

data_loader.py

@@ -0,0 +1,144 @@
+"""
+data_loader.py
+──────────────
+Handles all dataset loading, validation splitting, preprocessing and tokenisation.
+
+AG News label scheme:
+    0 = World   1 = Sports   2 = Business   3 = Sci/Tech
+"""
+import logging
+from typing import List, Optional, Tuple
+
+from datasets import load_dataset, DatasetDict
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from config import CFG
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s  %(levelname)-8s  %(message)s",
+    datefmt="%H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def load_ag_news(
+    max_train: Optional[int] = CFG.max_train_samples,
+    max_eval:  Optional[int] = CFG.max_eval_samples,
+    max_test:  Optional[int] = CFG.max_test_samples,
+) -> DatasetDict:
+    """
+    Load AG News from the HuggingFace datasets cache (downloads on first call).
+
+    AG News ships with 'train' (120 K) and 'test' (7.6 K) only.
+    We carve out a stratified 10 % of 'train' as the validation set.
+
+    Returns
+    -------
+    DatasetDict with splits: 'train', 'validation', 'test'
+    """
+    logger.info("Loading AG News dataset …")
+    raw = load_dataset("ag_news")
+
+    # Stratified 90/10 train → train + validation
+    tv = raw["train"].train_test_split(
+        test_size=0.10,
+        seed=CFG.seed,
+        stratify_by_column="label",
+    )
+    dataset = DatasetDict({
+        "train":      tv["train"],
+        "validation": tv["test"],
+        "test":       raw["test"],
+    })
+
+    # Optional down-sampling (speeds up CPU training significantly)
+    if max_train is not None:
+        n = min(max_train, len(dataset["train"]))
+        dataset["train"] = (
+            dataset["train"].shuffle(seed=CFG.seed).select(range(n))
+        )
+    if max_eval is not None:
+        n = min(max_eval, len(dataset["validation"]))
+        dataset["validation"] = (
+            dataset["validation"].shuffle(seed=CFG.seed).select(range(n))
+        )
+    if max_test is not None:
+        n = min(max_test, len(dataset["test"]))
+        dataset["test"] = dataset["test"].select(range(n))
+
+    logger.info(
+        f"  train={len(dataset['train']):,}  "
+        f"val={len(dataset['validation']):,}  "
+        f"test={len(dataset['test']):,}"
+    )
+    return dataset
+
+
+def load_test_only() -> Tuple[List[str], List[int]]:
+    """
+    Load only the test split (fast, no stratified split overhead).
+    Used by compare_results.py.
+    """
+    raw = load_dataset("ag_news")
+    return list(raw["test"]["text"]), list(raw["test"]["label"])
+
+
+def get_raw_splits(dataset: DatasetDict) -> Tuple:
+    """
+    Return plain Python lists of (texts, labels) for all three splits.
+    Used by the scikit-learn traditional ML pipeline.
+    """
+    X_train = list(dataset["train"]["text"])
+    y_train = list(dataset["train"]["label"])
+    X_val   = list(dataset["validation"]["text"])
+    y_val   = list(dataset["validation"]["label"])
+    X_test  = list(dataset["test"]["text"])
+    y_test  = list(dataset["test"]["label"])
+    return X_train, y_train, X_val, y_val, X_test, y_test
+
+
+def get_tokenizer() -> PreTrainedTokenizerBase:
+    """Download (or load from local HuggingFace cache) the DistilBERT tokeniser."""
+    logger.info(f"Loading tokeniser: {CFG.model_checkpoint}")
+    return AutoTokenizer.from_pretrained(CFG.model_checkpoint)
+
+
+def tokenise_dataset(
+    dataset: DatasetDict,
+    tokenizer: PreTrainedTokenizerBase,
+) -> DatasetDict:
+    """
+    Tokenise all splits for the HuggingFace Trainer.
+
+    Design decisions:
+    - padding=False  → pads at collation time via DataCollatorWithPadding
+                       (more memory-efficient than padding all to max_length)
+    - num_proc=1     → required on Windows; fork-based multi-processing
+                       causes issues with PyTorch on Windows
+    """
+    def _tokenise(batch: dict) -> dict:
+        return tokenizer(
+            batch["text"],
+            truncation=True,
+            max_length=CFG.max_length,
+            padding=False,
+        )
+
+    logger.info("Tokenising dataset …")
+    tokenised = dataset.map(
+        _tokenise,
+        batched=True,
+        batch_size=1_000,
+        num_proc=1,
+        remove_columns=["text"],
+        desc="Tokenising",
+    )
+
+    # HuggingFace Trainer requires the label column to be named 'labels'
+    tokenised = tokenised.rename_column("label", "labels")
+    tokenised.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
+
+    return tokenised

database.py

@@ -0,0 +1,383 @@
+import os
+import sqlite3
+import threading
+import uuid
+from datetime import datetime, timedelta, timezone
+from typing import Any, Dict, List, Optional, Tuple
+
+from config import CFG
+
+
+_WRITE_LOCK = threading.Lock()
+
+
+def _logs_dir() -> str:
+    path = os.path.join("logs")
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
+def _default_db_path() -> str:
+    return os.path.join(_logs_dir(), "api_requests.db")
+
+
+def _connect(db_path: Optional[str] = None) -> sqlite3.Connection:
+    conn = sqlite3.connect(db_path or _default_db_path(), timeout=30, check_same_thread=False)
+    conn.row_factory = sqlite3.Row
+    conn.execute("PRAGMA journal_mode=WAL;")
+    conn.execute("PRAGMA synchronous=NORMAL;")
+    conn.execute("PRAGMA foreign_keys=ON;")
+    return conn
+
+
+def _now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _today_ymd() -> str:
+    return datetime.now(timezone.utc).date().isoformat()
+
+
+def init_db(db_path: Optional[str] = None) -> None:
+    with _WRITE_LOCK:
+        conn = _connect(db_path=db_path)
+        try:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS requests (
+                  id INTEGER PRIMARY KEY AUTOINCREMENT,
+                  request_id TEXT UNIQUE NOT NULL,
+                  timestamp TEXT NOT NULL,
+                  model_name TEXT NOT NULL,
+                  input_text TEXT NOT NULL,
+                  input_length INTEGER,
+                  predicted_label TEXT NOT NULL,
+                  predicted_label_id INTEGER NOT NULL,
+                  confidence REAL NOT NULL,
+                  is_low_confidence INTEGER NOT NULL DEFAULT 0,
+                  latency_ms REAL NOT NULL,
+                  is_batch INTEGER NOT NULL DEFAULT 0
+                );
+                """
+            )
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS model_stats (
+                  model_name TEXT NOT NULL,
+                  date TEXT NOT NULL,
+                  total_requests INTEGER DEFAULT 0,
+                  avg_confidence REAL DEFAULT 0.0,
+                  avg_latency_ms REAL DEFAULT 0.0,
+                  low_conf_count INTEGER DEFAULT 0,
+                  PRIMARY KEY (model_name, date)
+                );
+                """
+            )
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS low_confidence_flags (
+                  id INTEGER PRIMARY KEY AUTOINCREMENT,
+                  request_id TEXT NOT NULL,
+                  timestamp TEXT NOT NULL,
+                  input_text TEXT NOT NULL,
+                  predicted_label TEXT NOT NULL,
+                  confidence REAL NOT NULL,
+                  reviewed INTEGER NOT NULL DEFAULT 0,
+                  review_note TEXT,
+                  FOREIGN KEY (request_id) REFERENCES requests(request_id)
+                );
+                """
+            )
+            conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_requests_timestamp ON requests(timestamp);"
+            )
+            conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_requests_model ON requests(model_name);"
+            )
+            conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_flags_reviewed ON low_confidence_flags(reviewed);"
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+
+def new_request_id() -> str:
+    return str(uuid.uuid4())
+
+
+def log_request(
+    request_id: str,
+    model_name: str,
+    input_text: str,
+    predicted_label: str,
+    predicted_label_id: int,
+    confidence: float,
+    latency_ms: float,
+    is_batch: bool,
+    db_path: Optional[str] = None,
+) -> None:
+    ts = _now_iso()
+    original_len = len(input_text)
+    stored_text = input_text[:500]
+    is_low = 1 if float(confidence) < float(CFG.low_confidence_threshold) else 0
+    batch_int = 1 if is_batch else 0
+
+    with _WRITE_LOCK:
+        conn = _connect(db_path=db_path)
+        try:
+            conn.execute(
+                """
+                INSERT INTO requests (
+                  request_id, timestamp, model_name, input_text, input_length,
+                  predicted_label, predicted_label_id, confidence, is_low_confidence,
+                  latency_ms, is_batch
+                )
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
+                """,
+                (
+                    request_id,
+                    ts,
+                    model_name,
+                    stored_text,
+                    original_len,
+                    predicted_label,
+                    int(predicted_label_id),
+                    float(confidence),
+                    int(is_low),
+                    float(latency_ms),
+                    int(batch_int),
+                ),
+            )
+
+            if is_low:
+                conn.execute(
+                    """
+                    INSERT INTO low_confidence_flags (
+                      request_id, timestamp, input_text, predicted_label, confidence, reviewed, review_note
+                    )
+                    VALUES (?, ?, ?, ?, ?, 0, NULL);
+                    """,
+                    (request_id, ts, stored_text, predicted_label, float(confidence)),
+                )
+
+            date = _today_ymd()
+            row = conn.execute(
+                """
+                SELECT total_requests, avg_confidence, avg_latency_ms, low_conf_count
+                FROM model_stats
+                WHERE model_name=? AND date=?;
+                """,
+                (model_name, date),
+            ).fetchone()
+            if row is None:
+                conn.execute(
+                    """
+                    INSERT INTO model_stats (
+                      model_name, date, total_requests, avg_confidence, avg_latency_ms, low_conf_count
+                    )
+                    VALUES (?, ?, 1, ?, ?, ?);
+                    """,
+                    (model_name, date, float(confidence), float(latency_ms), int(is_low)),
+                )
+            else:
+                n = int(row["total_requests"])
+                new_n = n + 1
+                new_avg_conf = (float(row["avg_confidence"]) * n + float(confidence)) / new_n
+                new_avg_lat = (float(row["avg_latency_ms"]) * n + float(latency_ms)) / new_n
+                new_low = int(row["low_conf_count"]) + int(is_low)
+                conn.execute(
+                    """
+                    UPDATE model_stats
+                    SET total_requests=?, avg_confidence=?, avg_latency_ms=?, low_conf_count=?
+                    WHERE model_name=? AND date=?;
+                    """,
+                    (new_n, new_avg_conf, new_avg_lat, new_low, model_name, date),
+                )
+
+            conn.commit()
+        finally:
+            conn.close()
+
+
+def get_request_history(
+    db_path: Optional[str] = None, limit: int = 100, offset: int = 0
+) -> List[Dict[str, Any]]:
+    conn = _connect(db_path=db_path)
+    try:
+        rows = conn.execute(
+            """
+            SELECT *
+            FROM requests
+            ORDER BY id DESC
+            LIMIT ? OFFSET ?;
+            """,
+            (int(limit), int(offset)),
+        ).fetchall()
+        return [dict(r) for r in rows]
+    finally:
+        conn.close()
+
+
+def get_low_confidence_flags(
+    db_path: Optional[str] = None, reviewed: bool = False, limit: int = 50
+) -> List[Dict[str, Any]]:
+    conn = _connect(db_path=db_path)
+    try:
+        rows = conn.execute(
+            """
+            SELECT *
+            FROM low_confidence_flags
+            WHERE reviewed=?
+            ORDER BY id DESC
+            LIMIT ?;
+            """,
+            (1 if reviewed else 0, int(limit)),
+        ).fetchall()
+        return [dict(r) for r in rows]
+    finally:
+        conn.close()
+
+
+def mark_reviewed(request_id: str, note: Optional[str] = None, db_path: Optional[str] = None) -> None:
+    with _WRITE_LOCK:
+        conn = _connect(db_path=db_path)
+        try:
+            conn.execute(
+                """
+                UPDATE low_confidence_flags
+                SET reviewed=1, review_note=?
+                WHERE request_id=?;
+                """,
+                (note, request_id),
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+
+def get_model_leaderboard(db_path: Optional[str] = None) -> List[Tuple[str, int, float, float]]:
+    conn = _connect(db_path=db_path)
+    try:
+        rows = conn.execute(
+            """
+            SELECT
+              model_name,
+              COUNT(*) AS total_requests,
+              AVG(confidence) AS avg_confidence,
+              AVG(latency_ms) AS avg_latency_ms
+            FROM requests
+            GROUP BY model_name
+            ORDER BY total_requests DESC;
+            """
+        ).fetchall()
+        return [
+            (
+                str(r["model_name"]),
+                int(r["total_requests"]),
+                float(r["avg_confidence"] or 0.0),
+                float(r["avg_latency_ms"] or 0.0),
+            )
+            for r in rows
+        ]
+    finally:
+        conn.close()
+
+
+def get_summary(
+    db_path: Optional[str] = None, model_name: Optional[str] = None, days: int = 7
+) -> Dict[str, Any]:
+    conn = _connect(db_path=db_path)
+    try:
+        start_ts = (datetime.now(timezone.utc) - timedelta(days=int(days))).isoformat()
+        params: List[Any] = [start_ts]
+        where = "WHERE timestamp >= ?"
+        if model_name:
+            where += " AND model_name = ?"
+            params.append(model_name)
+
+        row = conn.execute(
+            f"""
+            SELECT
+              COUNT(*) AS total_requests,
+              AVG(confidence) AS avg_confidence,
+              AVG(latency_ms) AS avg_latency_ms,
+              SUM(is_low_confidence) AS low_confidence_count
+            FROM requests
+            {where};
+            """,
+            tuple(params),
+        ).fetchone()
+
+        total_requests = int(row["total_requests"] or 0)
+        avg_confidence = float(row["avg_confidence"] or 0.0)
+        avg_latency_ms = float(row["avg_latency_ms"] or 0.0)
+        low_conf_count = int(row["low_confidence_count"] or 0)
+        rate = (low_conf_count / total_requests) * 100.0 if total_requests > 0 else 0.0
+
+        params2: List[Any] = list(params)
+        where2 = where
+
+        models = conn.execute(
+            f"""
+            SELECT DISTINCT model_name
+            FROM requests
+            {where2}
+            ORDER BY model_name;
+            """,
+            tuple(params2),
+        ).fetchall()
+        models_used = [str(r["model_name"]) for r in models]
+
+        label_rows = conn.execute(
+            f"""
+            SELECT predicted_label, COUNT(*) AS c
+            FROM requests
+            {where2}
+            GROUP BY predicted_label;
+            """,
+            tuple(params2),
+        ).fetchall()
+        predictions_by_label = {str(r["predicted_label"]): int(r["c"]) for r in label_rows}
+
+        return {
+            "period_days": int(days),
+            "total_requests": total_requests,
+            "models_used": models_used,
+            "avg_confidence": round(avg_confidence, 3),
+            "avg_latency_ms": round(avg_latency_ms, 2),
+            "low_confidence_count": low_conf_count,
+            "low_confidence_rate": f"{rate:.2f}%",
+            "predictions_by_label": predictions_by_label,
+        }
+    finally:
+        conn.close()
+
+
+def export_low_confidence_to_folder(
+    output_dir: str = os.path.join("logs", "low_confidence_review"),
+) -> Dict[str, Any]:
+    os.makedirs(output_dir, exist_ok=True)
+    flags = get_low_confidence_flags(reviewed=False, limit=10_000)
+    exported = 0
+    for f in flags:
+        request_id = str(f["request_id"])
+        ts = str(f["timestamp"]).replace(":", "-")
+        filename = f"{ts}_{request_id}.txt"
+        path = os.path.join(output_dir, filename)
+        if os.path.exists(path):
+            continue
+        content = "\n".join(
+            [
+                f"request_id: {request_id}",
+                f"timestamp: {f['timestamp']}",
+                f"predicted_label: {f['predicted_label']}",
+                f"confidence: {float(f['confidence']):.4f}",
+                "",
+                str(f["input_text"]),
+            ]
+        )
+        with open(path, "w", encoding="utf-8") as out:
+            out.write(content)
+        exported += 1
+    return {"exported": exported, "folder": output_dir}

download_dataset.py

@@ -0,0 +1,49 @@
+r"""
+download_dataset.py
+───────────────────
+One-time script to download and cache the AG News dataset from HuggingFace Hub.
+
+Run this BEFORE training to ensure the dataset is fully cached locally.
+The dataset is ~30 MB and is stored in:
+    Windows: C:\Users\<you>\.cache\huggingface\datasets\ag_news\
+
+Usage
+─────
+    python download_dataset.py
+"""
+import logging
+from datasets import load_dataset
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s  %(levelname)s  %(message)s",
+    datefmt="%H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+
+def main() -> None:
+    print("\n" + "─" * 60)
+    print("  Downloading AG News dataset from HuggingFace Hub")
+    print("  Size: ~30 MB  |  Cached after first download")
+    print("─" * 60 + "\n")
+
+    ds = load_dataset("ag_news")
+
+    print("\n  ✓  Download complete!")
+    print(f"\n     Split     Count")
+    print(f"     ──────    ──────────")
+    print(f"     train     {len(ds['train']):>10,}")
+    print(f"     test      {len(ds['test']):>10,}")
+    print(f"\n     Labels: 0=World  1=Sports  2=Business  3=Sci/Tech\n")
+
+    # Show one sample to confirm it loaded correctly
+    sample = ds["train"][42]
+    print("  Sample record (index 42):")
+    print(f"    text  : {sample['text'][:110]}…")
+    print(f"    label : {sample['label']}")
+    print()
+
+
+if __name__ == "__main__":
+    main()

ensemble.py

@@ -0,0 +1,349 @@
+"""
+ensemble.py
+-----------
+Soft-voting ensemble that combines all trained classifiers.
+Each model's class probabilities are weighted and summed for a final prediction.
+
+Usage
+-----
+    # Interactive predictions
+    python ensemble.py --interactive
+
+    # Single prediction
+    python ensemble.py --text "Tesla stock hits all-time high after earnings beat"
+
+    # Custom weights (must sum to 1.0)
+    python ensemble.py --text "..." --weights 0.05 0.10 0.85
+
+    # Use optimised weights from optimal_weights.json
+    python ensemble.py --text "..." --optimal
+"""
+import argparse
+import json
+import logging
+import os
+import sys
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+
+from config import CFG
+import traditional_model as tm
+import transformer_model as trm
+
+logging.basicConfig(level=logging.WARNING)
+
+# Path where optimize_ensemble.py saves the best weights
+_OPTIMAL_WEIGHTS_FILE = os.path.join(
+    CFG.outputs_dir, "ensemble_cache", "optimal_weights.json"
+)
+
+# Default model names used in this ensemble
+_DEFAULT_MODELS = ["lr", "svm", "distilbert_base_uncased"]
+_DEFAULT_WEIGHTS = [0.10, 0.15, 0.75]
+
+
+# -- Probability helpers ------------------------------------------------------
+
+def _proba_sklearn(text: str, pipeline) -> np.ndarray:
+    clf = list(pipeline.named_steps.values())[-1]
+    if hasattr(clf, "predict_proba"):
+        return pipeline.predict_proba([text])[0]
+    # LinearSVC: pseudo-probabilities via softmax over decision scores
+    scores = pipeline.decision_function([text])[0]
+    scores -= scores.max()
+    exp = np.exp(scores)
+    return exp / exp.sum()
+
+
+def _proba_transformer(text: str, model, tokenizer) -> np.ndarray:
+    enc = tokenizer(
+        text,
+        truncation=True,
+        max_length=CFG.max_length,
+        return_tensors="pt",
+    )
+    with torch.no_grad():
+        logits = model(**enc).logits[0]
+    return torch.softmax(logits, dim=-1).numpy()
+
+
+# -- Optimal weights loader ---------------------------------------------------
+
+def load_optimal_weights(
+    model_names: List[str],
+) -> Optional[Dict[str, float]]:
+    """
+    Attempt to load optimised weights from optimal_weights.json.
+
+    Returns a dict mapping model_name -> weight, or None if the file is
+    missing or malformed.
+    """
+    if not os.path.exists(_OPTIMAL_WEIGHTS_FILE):
+        logging.warning(
+            f"[Ensemble] Optimal weights file not found at "
+            f"'{_OPTIMAL_WEIGHTS_FILE}'. "
+            f"Run: python optimize_ensemble.py"
+        )
+        return None
+    try:
+        with open(_OPTIMAL_WEIGHTS_FILE) as fh:
+            data = json.load(fh)
+        weights = {name: data[name] for name in model_names if name in data}
+        if len(weights) != len(model_names):
+            logging.warning(
+                "[Ensemble] optimal_weights.json does not contain weights "
+                "for all requested models. Falling back to manual weights."
+            )
+            return None
+        logging.info(
+            f"[Ensemble] Loaded optimal weights (method={data.get('method')}, "
+            f"val_f1={data.get('val_f1_macro')}): {weights}"
+        )
+        return weights
+    except Exception as exc:
+        logging.warning(
+            f"[Ensemble] Could not load optimal_weights.json: {exc}. "
+            f"Falling back to manual weights."
+        )
+        return None
+
+
+# -- Ensemble class -----------------------------------------------------------
+
+class Ensemble:
+    """
+    Weighted soft-voting ensemble.
+
+    Parameters
+    ----------
+    model_weights : list of (model_name, weight) tuples.
+                    Weights are normalised automatically.
+                    model_name must match a key in saved_models/
+                    ('lr', 'svm', 'distilbert_base_uncased', etc.)
+    use_optimal_weights : bool, default True
+                    If True, attempt to load weights from
+                    outputs/ensemble_cache/optimal_weights.json and
+                    override the provided model_weights.
+                    Falls back to the provided weights if the file is
+                    missing or malformed.
+
+    Example
+    -------
+    >>> e = Ensemble([("lr", 0.10), ("svm", 0.15), ("distilbert_base_uncased", 0.75)])
+    >>> e.predict("Apple M5 chip breaks all benchmarks")
+
+    >>> # Load with auto-optimised weights
+    >>> e = Ensemble.from_optimal()
+    """
+
+    def __init__(
+        self,
+        model_weights: List[Tuple[str, float]],
+        use_optimal_weights: bool = True,
+    ):
+        # Attempt to override with optimised weights
+        if use_optimal_weights:
+            names = [name for name, _ in model_weights]
+            optimal = load_optimal_weights(names)
+            if optimal is not None:
+                model_weights = [(name, optimal[name]) for name in names]
+                print(
+                    f"  [Ensemble] Using optimal weights from "
+                    f"{_OPTIMAL_WEIGHTS_FILE}"
+                )
+
+        total = sum(w for _, w in model_weights)
+        self._weights: Dict[str, float] = {
+            name: w / total for name, w in model_weights
+        }
+        self._loaded: Dict = {}
+        self._kinds: Dict = {}
+        self._load_all()
+
+    # -- Class methods --------------------------------------------------------
+
+    @classmethod
+    def from_optimal(cls, fallback_weights: Optional[List[Tuple[str, float]]] = None):
+        """
+        Build an Ensemble using weights from optimal_weights.json.
+
+        If the file is missing, falls back to `fallback_weights` (or the
+        module-level defaults).
+
+        Parameters
+        ----------
+        fallback_weights : list of (model_name, weight) tuples, optional.
+                           Used when optimal_weights.json cannot be loaded.
+
+        Returns
+        -------
+        Ensemble instance
+        """
+        if fallback_weights is None:
+            fallback_weights = list(zip(_DEFAULT_MODELS, _DEFAULT_WEIGHTS))
+
+        # Try loading the optimal weights file directly
+        optimal = load_optimal_weights([name for name, _ in fallback_weights])
+        if optimal is not None:
+            weights = [(name, optimal[name]) for name, _ in fallback_weights]
+        else:
+            weights = fallback_weights
+
+        # Pass use_optimal_weights=False to avoid double-loading
+        return cls(weights, use_optimal_weights=False)
+
+    # -- Internal helpers -----------------------------------------------------
+
+    def _load_all(self) -> None:
+        for name in self._weights:
+            print(f"  Loading: {name} ...")
+            if name in ("lr", "svm"):
+                self._loaded[name] = tm.load_model(name)
+                self._kinds[name] = "sklearn"
+            else:
+                # Transformer: name is the directory under saved_models/
+                self._loaded[name] = trm.load_model(name)
+                self._kinds[name] = "transformer"
+        print()
+
+    def _proba(self, text: str, name: str) -> np.ndarray:
+        if self._kinds[name] == "sklearn":
+            return _proba_sklearn(text, self._loaded[name])
+        model, tokenizer = self._loaded[name]
+        return _proba_transformer(text, model, tokenizer)
+
+    # -- Public API -----------------------------------------------------------
+
+    def predict(self, text: str) -> Dict:
+        """
+        Compute the weighted ensemble prediction for a single text.
+
+        Returns predicted label, ensemble probabilities, and per-model
+        debug info.
+        """
+        combined = np.zeros(CFG.num_labels, dtype=float)
+        model_probs = {}
+
+        for name, weight in self._weights.items():
+            p = self._proba(text, name)
+            combined += weight * p
+            model_probs[name] = {
+                CFG.label_names[i]: round(float(p[i]), 4)
+                for i in range(CFG.num_labels)
+            }
+
+        pred_id = int(np.argmax(combined))
+        return {
+            "text": text,
+            "label_id": pred_id,
+            "label": CFG.label_names[pred_id],
+            "confidence": round(float(combined[pred_id]), 4),
+            "ensemble_probabilities": {
+                CFG.label_names[i]: round(float(combined[i]), 4)
+                for i in range(CFG.num_labels)
+            },
+            "per_model": model_probs,
+        }
+
+    @property
+    def weights(self) -> Dict[str, float]:
+        """Return the normalised per-model weights."""
+        return dict(self._weights)
+
+
+# -- Display ------------------------------------------------------------------
+
+def display(result: Dict) -> None:
+    snippet = result["text"][:88] + ("..." if len(result["text"]) > 88 else "")
+    print(f"\n  Input      : {snippet}")
+    print(f"  Label      : [{result['label_id']}]  {result['label']}")
+    print(f"  Confidence : {result['confidence']:.4f}")
+    print("  Ensemble Scores:")
+    for label, prob in sorted(
+        result["ensemble_probabilities"].items(),
+        key=lambda x: x[1],
+        reverse=True,
+    ):
+        bar = "#" * round(prob * 28)
+        print(f"    {label:<12}  [{bar:<28}]  {prob:.4f}")
+    print()
+
+
+# -- CLI ----------------------------------------------------------------------
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Ensemble Document Classifier")
+    parser.add_argument(
+        "--text", type=str, default=None, help="Single text to classify"
+    )
+    parser.add_argument(
+        "--interactive",
+        action="store_true",
+        help="Enter interactive prediction loop",
+    )
+    parser.add_argument(
+        "--weights",
+        nargs=3,
+        type=float,
+        default=_DEFAULT_WEIGHTS,
+        metavar=("LR_W", "SVM_W", "DISTILBERT_W"),
+        help="Weights for LR, SVM, DistilBERT (auto-normalised)",
+    )
+    parser.add_argument(
+        "--optimal",
+        action="store_true",
+        default=False,
+        help="Load weights from optimal_weights.json (ignores --weights)",
+    )
+    parser.add_argument(
+        "--no-optimal",
+        dest="optimal",
+        action="store_false",
+        help="Disable automatic loading of optimal weights",
+    )
+    args = parser.parse_args()
+
+    print("\n  Building Ensemble ...")
+    model_weights = [
+        ("lr",                      args.weights[0]),
+        ("svm",                     args.weights[1]),
+        ("distilbert_base_uncased", args.weights[2]),
+    ]
+
+    # --optimal flag forces loading optimal weights; otherwise honour
+    # use_optimal_weights=True default (auto-load if file exists)
+    use_optimal = True  # always attempt; falls back gracefully
+    if args.optimal:
+        ensemble = Ensemble.from_optimal(fallback_weights=model_weights)
+    else:
+        ensemble = Ensemble(model_weights, use_optimal_weights=use_optimal)
+
+    print(f"  Ensemble ready.  Active weights: {ensemble.weights}\n")
+
+    if args.interactive:
+        print("  Ensemble -- Interactive Mode  |  Type 'q' to exit\n")
+        while True:
+            try:
+                text = input("  >> ").strip()
+            except (KeyboardInterrupt, EOFError):
+                print("\n  Bye.")
+                break
+            if not text:
+                continue
+            if text.lower() in {"q", "quit", "exit"}:
+                print("  Bye.")
+                break
+            display(ensemble.predict(text))
+
+    elif args.text:
+        display(ensemble.predict(args.text))
+
+    else:
+        parser.print_help()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

error_analysis.py

@@ -0,0 +1,211 @@
+"""
+error_analysis.py
+─────────────────
+Detailed analysis of model errors on the test set.
+Generates confidence distributions, per-class accuracy bars,
+and a CSV of the hardest misclassified examples.
+
+Usage
+─────
+    python error_analysis.py --model roberta-base
+    python error_analysis.py --model lr
+    python error_analysis.py --model svm
+"""
+import argparse
+import logging
+import os
+from typing import List, Tuple
+
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+import torch
+from sklearn.metrics import accuracy_score
+
+from config import CFG
+from data_loader import load_test_only
+import traditional_model as tm
+import transformer_model as trm
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s  %(levelname)-8s  %(message)s",
+    datefmt="%H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+
+# ── Probability extraction ─────────────────────────────────────────────────────
+
+def _proba_sklearn(text_list: List[str], pipeline) -> np.ndarray:
+    clf = list(pipeline.named_steps.values())[-1]
+    if hasattr(clf, "predict_proba"):
+        return pipeline.predict_proba(text_list)
+    # LinearSVC: convert decision scores to pseudo-probabilities via softmax
+    scores = pipeline.decision_function(text_list)
+    scores -= scores.max(axis=1, keepdims=True)
+    exp     = np.exp(scores)
+    return exp / exp.sum(axis=1, keepdims=True)
+
+
+def _proba_transformer(text_list: List[str], model, tokenizer) -> np.ndarray:
+    all_probs  = []
+    batch_size = 32
+    for i in range(0, len(text_list), batch_size):
+        batch = text_list[i : i + batch_size]
+        enc   = tokenizer(batch, truncation=True, max_length=CFG.max_length,
+                          padding=True, return_tensors="pt")
+        with torch.no_grad():
+            logits = model(**enc).logits
+        all_probs.append(torch.softmax(logits, dim=-1).numpy())
+    return np.vstack(all_probs)
+
+
+# ── Main analysis ─────────────────────────────────────────────────────────────
+
+def analyse(model_name: str, save_dir: str = None) -> pd.DataFrame:
+    """
+    Full error analysis pipeline.
+
+    Returns
+    -------
+    DataFrame of all misclassified examples.
+    """
+    logger.info("Loading test set …")
+    X_test, y_test = load_test_only()
+
+    logger.info(f"Running predictions with: {model_name}")
+    if model_name in ("lr", "svm"):
+        pipeline = tm.load_model(model_name)
+        proba    = _proba_sklearn(X_test, pipeline)
+        preds    = proba.argmax(axis=1).tolist()
+    else:
+        model, tokenizer = trm.load_model(model_name)
+        proba = _proba_transformer(X_test, model, tokenizer)
+        preds = proba.argmax(axis=1).tolist()
+
+    acc = accuracy_score(y_test, preds)
+    logger.info(f"Test accuracy: {acc * 100:.2f}%")
+
+    # Build analysis DataFrame
+    df = pd.DataFrame({
+        "text":       X_test,
+        "true_label": [CFG.label_names[y] for y in y_test],
+        "pred_label": [CFG.label_names[p] for p in preds],
+        "confidence": proba.max(axis=1),
+        "correct":    [int(y) == int(p) for y, p in zip(y_test, preds)],
+    })
+    for i, name in enumerate(CFG.label_names):
+        df[f"prob_{name}"] = proba[:, i]
+
+    errors   = df[~df["correct"].astype(bool)]
+    corrects = df[df["correct"].astype(bool)]
+
+    # ── Console report ───────────────────────────────────────────────────────
+    print("\n" + "═" * 60)
+    print(f"  ERROR ANALYSIS  —  {model_name.upper()}")
+    print("═" * 60)
+    print(f"  Total   : {len(df):,}")
+    print(f"  Correct : {len(corrects):,}  ({len(corrects)/len(df)*100:.2f}%)")
+    print(f"  Errors  : {len(errors):,}   ({len(errors)/len(df)*100:.2f}%)")
+
+    print("\n  Errors by true class:")
+    for label in CFG.label_names:
+        n = len(errors[errors["true_label"] == label])
+        print(f"    {label:<12}  {n:>4} errors")
+
+    print("\n  Top confused pairs  (True → Predicted):")
+    confused = (
+        errors.groupby(["true_label", "pred_label"])
+        .size()
+        .sort_values(ascending=False)
+        .head(6)
+    )
+    for (true, pred), count in confused.items():
+        print(f"    {true:<12} → {pred:<12}  {count:>4} times")
+
+    print("\n  5 Hardest Errors (lowest confidence):")
+    for _, row in errors.nsmallest(5, "confidence").iterrows():
+        snippet = row["text"][:75] + "…"
+        print(f"    [{row['true_label']} → {row['pred_label']}  conf={row['confidence']:.3f}]")
+        print(f"    {snippet}\n")
+
+    # ── Plots ──────────────���─────────────────────────────────────────────────
+    _plot_analysis(df, model_name, save_dir)
+
+    # ── Save CSV ─────────────────────────────────────────────────────────────
+    if save_dir:
+        os.makedirs(save_dir, exist_ok=True)
+        csv_path = os.path.join(save_dir, f"errors_{model_name.replace('-','_')}.csv")
+        errors.to_csv(csv_path, index=False)
+        logger.info(f"Error CSV → {csv_path}")
+
+    return errors
+
+
+def _plot_analysis(df: pd.DataFrame, model_name: str, save_dir: str = None) -> None:
+    """Two-panel figure: confidence distribution + per-class accuracy bars."""
+    fig, axes = plt.subplots(1, 2, figsize=(13, 5))
+    fig.suptitle(f"Error Analysis — {model_name}", fontsize=14, fontweight="bold")
+
+    # Panel 1: Confidence histograms
+    correct_conf = df[df["correct"].astype(bool)]["confidence"]
+    error_conf   = df[~df["correct"].astype(bool)]["confidence"]
+    axes[0].hist(correct_conf, bins=30, alpha=0.75, color="#27ae60",
+                 label=f"Correct  (n={len(correct_conf):,})")
+    axes[0].hist(error_conf, bins=30, alpha=0.75, color="#e74c3c",
+                 label=f"Incorrect (n={len(error_conf):,})")
+    axes[0].set_xlabel("Prediction Confidence", fontsize=11)
+    axes[0].set_ylabel("Count", fontsize=11)
+    axes[0].set_title("Confidence Distribution", fontsize=12)
+    axes[0].legend(fontsize=10)
+    axes[0].axvline(correct_conf.mean(), color="#27ae60", linestyle="--", linewidth=1.2,
+                    label=f"Mean correct: {correct_conf.mean():.3f}")
+    axes[0].axvline(error_conf.mean(), color="#e74c3c", linestyle="--", linewidth=1.2,
+                    label=f"Mean error:   {error_conf.mean():.3f}")
+
+    # Panel 2: Per-class accuracy
+    colours     = ["#3498db", "#27ae60", "#e67e22", "#9b59b6"]
+    class_accs  = [
+        df[df["true_label"] == lbl]["correct"].astype(float).mean() * 100
+        for lbl in CFG.label_names
+    ]
+    bars = axes[1].bar(CFG.label_names, class_accs, color=colours,
+                       edgecolor="white", linewidth=1.5)
+    axes[1].set_ylim(80, 100)
+    axes[1].set_xlabel("Class", fontsize=11)
+    axes[1].set_ylabel("Accuracy (%)", fontsize=11)
+    axes[1].set_title("Per-Class Accuracy", fontsize=12)
+    for bar, acc in zip(bars, class_accs):
+        axes[1].text(bar.get_x() + bar.get_width() / 2,
+                     bar.get_height() + 0.3,
+                     f"{acc:.1f}%", ha="center", va="bottom", fontsize=11, fontweight="bold")
+
+    plt.tight_layout()
+
+    if save_dir:
+        os.makedirs(save_dir, exist_ok=True)
+        path = os.path.join(save_dir, f"analysis_{model_name.replace('-','_')}.png")
+        plt.savefig(path, dpi=150)
+        logger.info(f"Plot → {path}")
+
+    plt.show()
+    plt.close(fig)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Document classifier error analysis")
+    parser.add_argument(
+        "--model", default="roberta-base",
+        help="Model name: 'lr', 'svm', or transformer checkpoint (e.g. 'roberta-base')"
+    )
+    args = parser.parse_args()
+    save_dir = os.path.join(CFG.outputs_dir, "error_analysis")
+    analyse(args.model, save_dir=save_dir)
+
+
+if __name__ == "__main__":
+    main()

hyperparameter_search.py

@@ -0,0 +1,365 @@
+import argparse
+import json
+import os
+import warnings
+from pathlib import Path
+from typing import Any, Dict, Literal, Tuple
+ 
+import joblib
+import matplotlib.pyplot as plt
+import numpy as np
+import optuna
+import seaborn as sns
+from sklearn.exceptions import ConvergenceWarning
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
+from sklearn.pipeline import Pipeline
+from sklearn.svm import LinearSVC
+from tqdm import tqdm
+ 
+from config import CFG
+from data_loader import get_raw_splits, load_ag_news
+ 
+ModelType = Literal["lr", "svm"]
+ 
+BASELINE_TEST_ACCURACY: Dict[ModelType, float] = {"lr": 0.9045, "svm": 0.9089}
+DEFAULT_MAX_TRAIN = 40_000
+ 
+ 
+def _storage_url() -> str:
+    db_path = os.path.abspath(os.path.join(CFG.outputs_dir, "optuna_studies.db"))
+    return f"sqlite:///{Path(db_path).as_posix()}"
+ 
+ 
+def _hp_outputs_dir() -> str:
+    path = os.path.join(CFG.outputs_dir, "hp_search")
+    os.makedirs(path, exist_ok=True)
+    return path
+ 
+ 
+def _trial_params(trial: optuna.trial.Trial, model_type: ModelType) -> Dict[str, Any]:
+    params: Dict[str, Any] = {}
+    params["max_features"] = trial.suggest_int(
+        "max_features", 20_000, 100_000, step=10_000
+    )
+    params["ngram_min"] = trial.suggest_int("ngram_min", 1, 1)
+    params["ngram_max"] = trial.suggest_int("ngram_max", 1, 3)
+    params["min_df"] = trial.suggest_int("min_df", 1, 5)
+    params["sublinear_tf"] = trial.suggest_categorical("sublinear_tf", [True, False])
+ 
+    if model_type == "lr":
+        params["C"] = trial.suggest_float("C", 0.1, 20.0, log=True)
+        params["solver"] = trial.suggest_categorical("solver", ["saga", "lbfgs"])
+    else:
+        params["C"] = trial.suggest_float("C", 0.01, 10.0, log=True)
+ 
+    return params
+ 
+ 
+def _build_pipeline(model_type: ModelType, params: Dict[str, Any]) -> Pipeline:
+    tfidf = TfidfVectorizer(
+        max_features=int(params["max_features"]),
+        ngram_range=(int(params["ngram_min"]), int(params["ngram_max"])),
+        sublinear_tf=bool(params["sublinear_tf"]),
+        min_df=int(params["min_df"]),
+        strip_accents="unicode",
+        analyzer="word",
+        token_pattern=r"\w{1,}",
+        dtype=np.float32,
+    )
+ 
+    if model_type == "lr":
+        clf = LogisticRegression(
+            C=float(params["C"]),
+            solver=str(params["solver"]),
+            max_iter=2_000,
+            n_jobs=-1,
+            random_state=CFG.seed,
+        )
+    else:
+        clf = LinearSVC(
+            C=float(params["C"]),
+            max_iter=3_000,
+            random_state=CFG.seed,
+        )
+ 
+    return Pipeline([("tfidf", tfidf), (model_type, clf)])
+ 
+ 
+def _make_objective(
+    model_type: ModelType,
+    X_train,
+    y_train,
+    X_val,
+    y_val,
+):
+    def objective(trial: optuna.trial.Trial) -> float:
+        params = _trial_params(trial, model_type)
+        pipeline = _build_pipeline(model_type, params)
+ 
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=ConvergenceWarning)
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            pipeline.fit(X_train, y_train)
+ 
+        val_preds = pipeline.predict(X_val)
+        return float(f1_score(y_val, val_preds, average="macro"))
+ 
+    return objective
+ 
+ 
+def _save_confusion_matrix(
+    cm,
+    title: str,
+    save_path: str,
+) -> None:
+    fig, ax = plt.subplots(figsize=(7, 6))
+    sns.heatmap(
+        cm,
+        annot=True,
+        fmt="d",
+        cmap="Blues",
+        xticklabels=CFG.label_names,
+        yticklabels=CFG.label_names,
+        linewidths=0.5,
+        ax=ax,
+    )
+    ax.set_xlabel("Predicted Label", fontsize=11)
+    ax.set_ylabel("True Label", fontsize=11)
+    ax.set_title(title, fontsize=13, fontweight="bold")
+    plt.tight_layout()
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    plt.savefig(save_path, dpi=150)
+    plt.close(fig)
+ 
+ 
+def _save_optuna_plots(study: optuna.Study, model_type: ModelType) -> Tuple[str, str]:
+    out_dir = _hp_outputs_dir()
+    import optuna.visualization.matplotlib as ovm
+ 
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=optuna.exceptions.ExperimentalWarning)
+        warnings.filterwarnings("ignore", category=UserWarning)
+        ax1 = ovm.plot_parallel_coordinate(study)
+    fig1 = ax1.figure
+    fig1.tight_layout()
+    p1 = os.path.join(out_dir, f"{model_type}_parallel_coords.png")
+    fig1.savefig(p1, dpi=150)
+    plt.close(fig1)
+ 
+    p2 = os.path.join(out_dir, f"{model_type}_param_importance.png")
+    try:
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=optuna.exceptions.ExperimentalWarning)
+            warnings.filterwarnings("ignore", category=UserWarning)
+            ax2 = ovm.plot_param_importances(study)
+        fig2 = ax2.figure
+        fig2.tight_layout()
+        fig2.savefig(p2, dpi=150)
+        plt.close(fig2)
+    except ValueError:
+        p2 = ""
+ 
+    return p1, p2
+ 
+ 
+def _create_or_reset_study(
+    study_name: str,
+    storage: str,
+    resume: bool,
+) -> optuna.Study:
+    sampler = optuna.samplers.TPESampler(seed=CFG.seed)
+ 
+    if not resume:
+        try:
+            optuna.delete_study(study_name=study_name, storage=storage)
+        except KeyError:
+            pass
+ 
+    return optuna.create_study(
+        direction="maximize",
+        study_name=study_name,
+        storage=storage,
+        load_if_exists=True,
+        sampler=sampler,
+    )
+ 
+ 
+def _run_model_search(
+    model_type: ModelType,
+    n_trials_total: int,
+    full: bool,
+    resume: bool,
+) -> Dict[str, Any]:
+    max_train = None if full else DEFAULT_MAX_TRAIN
+ 
+    dataset = load_ag_news(max_train=max_train, max_eval=None, max_test=None)
+    X_train, y_train, X_val, y_val, X_test, y_test = get_raw_splits(dataset)
+ 
+    storage = _storage_url()
+    study_name = f"{model_type}_hyperparams"
+    study = _create_or_reset_study(study_name=study_name, storage=storage, resume=resume)
+ 
+    remaining = max(0, int(n_trials_total) - len(study.trials))
+    if remaining == 0:
+        best_params = study.best_params
+        best_val_f1 = float(study.best_value)
+    else:
+        objective = _make_objective(model_type, X_train, y_train, X_val, y_val)
+        pbar = tqdm(total=remaining, desc=f"{model_type.upper()} Optuna", unit="trial")
+ 
+        def _cb(_study: optuna.Study, _trial: optuna.trial.FrozenTrial) -> None:
+            pbar.update(1)
+ 
+        study.optimize(
+            objective,
+            n_trials=remaining,
+            callbacks=[_cb],
+            gc_after_trial=True,
+            show_progress_bar=False,
+        )
+        pbar.close()
+        best_params = study.best_params
+        best_val_f1 = float(study.best_value)
+ 
+    pipeline = _build_pipeline(model_type, best_params)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=ConvergenceWarning)
+        warnings.filterwarnings("ignore", category=FutureWarning)
+        if full:
+            X_final = list(X_train) + list(X_val)
+            y_final = list(y_train) + list(y_val)
+            pipeline.fit(X_final, y_final)
+        else:
+            pipeline.fit(X_train, y_train)
+ 
+    test_preds = pipeline.predict(X_test)
+    test_acc = float(accuracy_score(y_test, test_preds))
+    cm = confusion_matrix(y_test, test_preds)
+ 
+    out_dir = _hp_outputs_dir()
+    cm_path = os.path.join(out_dir, f"{model_type}_confusion_matrix.png")
+    _save_confusion_matrix(
+        cm,
+        title=f"Optimized {model_type.upper()} — Confusion Matrix",
+        save_path=cm_path,
+    )
+ 
+    model_path = os.path.join(CFG.models_dir, f"traditional_{model_type}_optimized.joblib")
+    joblib.dump(pipeline, model_path)
+ 
+    p1, p2 = _save_optuna_plots(study, model_type=model_type)
+ 
+    return {
+        "model": model_type,
+        "study_name": study_name,
+        "storage": storage,
+        "max_train": max_train,
+        "best_val_f1_macro": best_val_f1,
+        "best_params": best_params,
+        "test_accuracy": test_acc,
+        "confusion_matrix_path": cm_path,
+        "model_path": model_path,
+        "plot_parallel_coords_path": p1,
+        "plot_param_importance_path": p2,
+    }
+ 
+ 
+def _print_summary(results: Dict[ModelType, Dict[str, Any]]) -> None:
+    def _fmt_params(d: Dict[str, Any]) -> str:
+        s = json.dumps(d, sort_keys=True)
+        return s if len(s) <= 110 else s[:107] + "..."
+ 
+    headers = ["Model", "Best Val F1", "Best Params", "Improvement vs Baseline"]
+    rows = []
+    for model_type, r in results.items():
+        baseline = BASELINE_TEST_ACCURACY[model_type]
+        improvement_pp = (float(r["test_accuracy"]) - baseline) * 100.0
+        rows.append(
+            [
+                model_type.upper(),
+                f"{float(r['best_val_f1_macro']):.4f}",
+                _fmt_params(r["best_params"]),
+                f"{improvement_pp:+.2f} pp",
+            ]
+        )
+ 
+    col_widths = [
+        max(len(headers[i]), max(len(str(row[i])) for row in rows)) for i in range(4)
+    ]
+    fmt = " | ".join(f"{{:<{w}}}" for w in col_widths)
+    sep = "-+-".join("-" * w for w in col_widths)
+ 
+    print("\n" + fmt.format(*headers))
+    print(sep)
+    for row in rows:
+        print(fmt.format(*row))
+ 
+ 
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Optuna hyperparameter search for TF-IDF + LR/SVM.")
+    parser.add_argument(
+        "--model",
+        choices=["lr", "svm", "all"],
+        default="all",
+        help="Which model study to run.",
+    )
+    parser.add_argument(
+        "--n-trials",
+        type=int,
+        default=30,
+        help="Total trials per model study (respects --resume).",
+    )
+    parser.add_argument(
+        "--full",
+        action="store_true",
+        help="Use full 120K training examples (much slower per trial).",
+    )
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        help="Resume from the SQLite study DB (otherwise, resets the study).",
+    )
+    args = parser.parse_args()
+ 
+    optuna.logging.set_verbosity(optuna.logging.WARNING)
+ 
+    train_note = (
+        "FULL (120K)"
+        if args.full
+        else f"CAPPED ({DEFAULT_MAX_TRAIN:,} max_train for i3 CPU)"
+    )
+    print(
+        f"[HP Search] Train size: {train_note}. "
+        f"Override with --full to use the complete dataset."
+    )
+    print(f"[HP Search] Storage: {_storage_url()}")
+ 
+    results: Dict[ModelType, Dict[str, Any]] = {}
+    if args.model in ("lr", "all"):
+        results["lr"] = _run_model_search(
+            model_type="lr",
+            n_trials_total=args.n_trials,
+            full=args.full,
+            resume=args.resume,
+        )
+    if args.model in ("svm", "all"):
+        results["svm"] = _run_model_search(
+            model_type="svm",
+            n_trials_total=args.n_trials,
+            full=args.full,
+            resume=args.resume,
+        )
+ 
+    _print_summary(results)
+ 
+    out_dir = _hp_outputs_dir()
+    best_params_path = os.path.join(out_dir, "best_params.json")
+    with open(best_params_path, "w", encoding="utf-8") as f:
+        json.dump(results, f, indent=2)
+    print(f"\n[HP Search] Best params -> {best_params_path}")
+ 
+ 
+if __name__ == "__main__":
+    main()

predict.py

@@ -0,0 +1,267 @@
+"""
+predict.py
+──────────
+Run inference on new text with any trained classifier.
+
+Usage
+─────
+    # Single prediction — traditional models
+    python predict.py --model lr  --text "Federal Reserve cuts rates to near zero"
+    python predict.py --model svm --text "Ronaldo scores hat-trick in Champions League"
+
+    # Single prediction — transformer (FP32)
+    python predict.py --model transformer --text "NASA launches James Webb successor"
+
+    # Single prediction — INT8 quantized transformer (fast CPU inference)
+    python predict.py --model transformer_quantized --text "Federal Reserve raises rates"
+    python predict.py --model transformer_quantized --checkpoint roberta-base --text "..."
+
+    # Interactive loop
+    python predict.py --model transformer --interactive
+    python predict.py --model transformer_quantized --interactive
+    python predict.py --model lr --interactive
+"""
+import argparse
+import logging
+import sys
+from typing import Dict, Optional
+
+import numpy as np
+import torch
+
+from config import CFG
+import traditional_model as tm
+import transformer_model as trm
+
+# Suppress INFO logs during interactive prediction
+logging.basicConfig(level=logging.WARNING)
+
+
+# ── Prediction functions ──────────────────────────────────────────────────────
+
+def predict_traditional(text: str, model_name: str) -> Dict:
+    """Run a single prediction with a saved sklearn pipeline."""
+    pipeline = tm.load_model(model_name)
+    pred_id  = int(pipeline.predict([text])[0])
+
+    result: Dict = {
+        "text":     text,
+        "label_id": pred_id,
+        "label":    CFG.label_names[pred_id],
+    }
+
+    # Logistic Regression supports predict_proba; LinearSVC does not
+    clf = pipeline.named_steps[model_name]
+    if hasattr(clf, "predict_proba"):
+        probs = pipeline.predict_proba([text])[0]
+        result["probabilities"] = {
+            CFG.label_names[i]: round(float(p), 4)
+            for i, p in enumerate(probs)
+        }
+
+    return result
+
+
+def predict_transformer(
+    text: str,
+    model=None,
+    tokenizer=None,
+) -> Dict:
+    """Run a single prediction with a fine-tuned transformer (FP32, MPS/CPU)."""
+    if model is None or tokenizer is None:
+        model, tokenizer = trm.load_model()
+
+    encoding = tokenizer(
+        text,
+        truncation=True,
+        max_length=CFG.max_length,
+        return_tensors="pt",
+    )
+
+    with torch.no_grad():
+        logits = model(**encoding).logits[0]
+
+    probs   = torch.softmax(logits, dim=-1).numpy()
+    pred_id = int(np.argmax(probs))
+
+    return {
+        "text":     text,
+        "label_id": pred_id,
+        "label":    CFG.label_names[pred_id],
+        "probabilities": {
+            CFG.label_names[i]: round(float(p), 4)
+            for i, p in enumerate(probs)
+        },
+    }
+
+
+def predict_transformer_quantized(
+    text: str,
+    model=None,
+    tokenizer=None,
+    is_quantized: bool = True,
+) -> Dict:
+    """Run inference with the INT8 quantized model (CPU only)."""
+    if model is None or tokenizer is None:
+        raise ValueError("Pass a pre-loaded model and tokenizer.")
+
+    encoding = tokenizer(
+        text,
+        truncation=True,
+        max_length=CFG.max_length,
+        return_tensors="pt",
+    )
+    # INT8 quantized kernels only run on CPU
+    encoding = {k: v.to("cpu") for k, v in encoding.items()}
+
+    with torch.inference_mode():
+        logits = model(**encoding).logits[0]
+
+    probs   = torch.softmax(logits, dim=-1).numpy()
+    pred_id = int(np.argmax(probs))
+
+    label = "[INT8] " if is_quantized else "[FP32] "
+    return {
+        "text":        text,
+        "label_id":   pred_id,
+        "label":      CFG.label_names[pred_id],
+        "model_type": label.strip(),
+        "probabilities": {
+            CFG.label_names[i]: round(float(p), 4)
+            for i, p in enumerate(probs)
+        },
+    }
+
+
+# ── Display ───────────────────────────────────────────────────────────────────
+
+def display_result(result: Dict) -> None:
+    """Print a formatted prediction result to the terminal."""
+    snippet = result["text"]
+    if len(snippet) > 90:
+        snippet = snippet[:90] + "…"
+
+    model_tag = f"  [{result['model_type']}]" if "model_type" in result else ""
+    print(f"\n  Input  : {snippet}")
+    print(f"  Label  : [{result['label_id']}]  {result['label']}{model_tag}")
+
+    if "probabilities" in result:
+        print("  Scores :")
+        sorted_probs = sorted(
+            result["probabilities"].items(),
+            key=lambda x: x[1],
+            reverse=True,
+        )
+        for label, prob in sorted_probs:
+            bar   = "█" * round(prob * 28)
+            blank = " " * (28 - round(prob * 28))
+            print(f"    {label:<12}  [{bar}{blank}]  {prob:.4f}")
+    print()
+
+
+# ── CLI ───────────────────────────────────────────────────────────────────────
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        description="Document Classifier — Inference",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    p.add_argument(
+        "--model",
+        default="transformer",
+        help="Which saved model to load: lr, svm, transformer, transformer_quantized, or a specific variant like distilbert_quantized (default: transformer)",
+    )
+    p.add_argument(
+        "--checkpoint",
+        type=str,
+        default="distilbert-base-uncased",
+        help="HuggingFace checkpoint name for transformer/transformer_quantized "
+             "(default: distilbert-base-uncased)",
+    )
+    p.add_argument(
+        "--text",
+        type=str,
+        default=None,
+        help="Single text string to classify",
+    )
+    p.add_argument(
+        "--interactive",
+        action="store_true",
+        help="Enter an interactive prediction loop",
+    )
+    return p
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+
+    # Normalise model selection and automatically extract checkpoint if specified
+    model_lower = args.model.lower()
+    is_quant = "quant" in model_lower or "int8" in model_lower
+    
+    if "distilbert" in model_lower:
+        args.checkpoint = "distilbert-base-uncased"
+    elif "roberta" in model_lower:
+        args.checkpoint = "roberta-base"
+    elif "bert" in model_lower:
+        args.checkpoint = "bert-base-uncased"
+        
+    if is_quant:
+        args.model = "transformer_quantized"
+    elif args.model not in ["lr", "svm"]:
+        args.model = "transformer"
+
+    # Pre-load the model once (avoids reloading on every prediction in loops)
+    cached_model      = None
+    cached_tokenizer  = None
+    cached_quantized  = False
+
+    if args.model == "transformer":
+        print(f"  Loading transformer model ({args.checkpoint}) …")
+        cached_model, cached_tokenizer = trm.load_model(args.checkpoint)
+        print("  Model ready.\n")
+
+    elif args.model == "transformer_quantized":
+        print(f"  Loading quantized model ({args.checkpoint}) …")
+        cached_model, cached_tokenizer, cached_quantized = trm.load_quantized_model(
+            args.checkpoint
+        )
+        tag = "INT8 quantized" if cached_quantized else "FP32 (INT8 not found, fell back)"
+        print(f"  Model ready — {tag}.\n")
+
+    def _predict(text: str) -> Dict:
+        if args.model == "transformer":
+            return predict_transformer(text, cached_model, cached_tokenizer)
+        if args.model == "transformer_quantized":
+            return predict_transformer_quantized(
+                text, cached_model, cached_tokenizer, is_quantized=cached_quantized
+            )
+        return predict_traditional(text, args.model)
+
+    if args.interactive:
+        print("  Document Classifier — Interactive Mode")
+        print("  Type text and press Enter.  Type 'q' or 'quit' to exit.\n")
+        while True:
+            try:
+                text = input("  >> ").strip()
+            except (KeyboardInterrupt, EOFError):
+                print("\n  Exiting.")
+                break
+            if not text:
+                continue
+            if text.lower() in {"q", "quit", "exit"}:
+                print("  Goodbye.")
+                break
+            display_result(_predict(text))
+
+    elif args.text:
+        display_result(_predict(args.text))
+
+    else:
+        print("  Error: provide --text <string> or use --interactive.\n")
+        build_parser().print_help()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,25 @@
+[tool.black]
+line-length = 100
+target-version = ["py311"]
+exclude = '''
+/(
+  \.venv | frontend | outputs | data | saved_models | \.git
+)/
+'''
+
+[tool.isort]
+profile = "black"
+line_length = 100
+skip_glob = [".venv/*", "frontend/*"]
+
+[tool.flake8]
+max-line-length = 100
+extend-ignore = ["E203", "W503"]
+exclude = [".venv", "frontend", "__pycache__", "outputs", "data", "saved_models"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+markers = [
+  "slow: marks tests that load real datasets or ML models",
+]
+addopts = "-v --tb=short"

pytest.ini

@@ -0,0 +1,6 @@
+[pytest]
+testpaths = tests
+markers =
+    slow: marks tests that load real datasets (deselect with '-m "not slow"')
+addopts = -v --tb=short
+

quantize_model.py

@@ -0,0 +1,285 @@
+import argparse
+import json
+import os
+import shutil
+import time
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+from config import CFG
+from data_loader import load_test_only
+from transformer_model import _checkpoint_to_dir
+
+
+def _dir_size_bytes(path: str) -> int:
+    total = 0
+    for root, _, files in os.walk(path):
+        for f in files:
+            fp = os.path.join(root, f)
+            try:
+                total += os.path.getsize(fp)
+            except OSError:
+                pass
+    return total
+
+
+def _mb(n_bytes: int) -> float:
+    return float(n_bytes) / (1024.0 * 1024.0)
+
+
+def _copy_tokenizer_files(src_dir: str, dst_dir: str) -> List[str]:
+    os.makedirs(dst_dir, exist_ok=True)
+    whitelist = {
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "special_tokens_map.json",
+        "vocab.txt",
+        "merges.txt",
+        "added_tokens.json",
+        "sentencepiece.bpe.model",
+        "spiece.model",
+        "config.json",
+    }
+    copied: List[str] = []
+    for name in os.listdir(src_dir):
+        src = os.path.join(src_dir, name)
+        dst = os.path.join(dst_dir, name)
+        if not os.path.isfile(src):
+            continue
+        if name in whitelist or name.startswith("tokenizer"):
+            shutil.copy2(src, dst)
+            copied.append(name)
+    return copied
+
+
+def _load_fp32_model(fp32_dir: str):
+    model = AutoModelForSequenceClassification.from_pretrained(fp32_dir)
+    tokenizer = AutoTokenizer.from_pretrained(fp32_dir)
+    model.eval()
+    model.to("cpu")
+    return model, tokenizer
+
+
+def _quantize_dynamic_int8(model_fp32: torch.nn.Module) -> torch.nn.Module:
+    # Apple Silicon (ARM) requires qnnpack; x86 defaults to fbgemm which is unavailable on MPS.
+    torch.backends.quantized.engine = "qnnpack"
+    model_int8 = torch.quantization.quantize_dynamic(
+        model_fp32,
+        {torch.nn.Linear},
+        dtype=torch.qint8,
+    )
+    model_int8.eval()
+    return model_int8
+
+
+def _batched(iterable: List[str], batch_size: int):
+    for i in range(0, len(iterable), batch_size):
+        yield iterable[i : i + batch_size]
+
+
+def _predict(
+    model: torch.nn.Module,
+    tokenizer,
+    texts: List[str],
+    batch_size: int,
+) -> np.ndarray:
+    preds: List[int] = []
+    with torch.inference_mode():
+        for batch in _batched(texts, batch_size):
+            enc = tokenizer(
+                batch,
+                truncation=True,
+                max_length=CFG.max_length,
+                padding=True,
+                return_tensors="pt",
+            )
+            enc = {k: v.to("cpu") for k, v in enc.items()}
+            logits = model(**enc).logits
+            batch_preds = torch.argmax(logits, dim=-1).cpu().numpy().tolist()
+            preds.extend(batch_preds)
+    return np.asarray(preds, dtype=np.int64)
+
+
+def _accuracy(
+    model: torch.nn.Module,
+    tokenizer,
+    X_test: List[str],
+    y_test: List[int],
+    batch_size: int = 32,
+) -> float:
+    y_pred = _predict(model, tokenizer, X_test, batch_size=batch_size)
+    y_true = np.asarray(y_test, dtype=np.int64)
+    return float((y_pred == y_true).mean())
+
+
+def _benchmark_latency_ms(
+    model: torch.nn.Module,
+    tokenizer,
+    sample_texts: List[str],
+    batch_size: int,
+    runs: int = 50,
+    warmup: int = 5,
+) -> float:
+    per_text_ms: List[float] = []
+    for i in range(runs):
+        t0 = time.perf_counter()
+        _predict(model, tokenizer, sample_texts, batch_size=batch_size)
+        dt = time.perf_counter() - t0
+        if i >= warmup:
+            per_text_ms.append((dt / len(sample_texts)) * 1000.0)
+    return float(np.median(per_text_ms))
+
+
+def _save_quantized_model(
+    model_int8: torch.nn.Module,
+    fp32_dir: str,
+    int8_dir: str,
+    checkpoint_dir_name: str,
+) -> Dict:
+    os.makedirs(int8_dir, exist_ok=True)
+    model_path = os.path.join(int8_dir, "model_int8.pt")
+    torch.save(model_int8, model_path)
+    _copy_tokenizer_files(fp32_dir, int8_dir)
+
+    original_size_mb = _mb(_dir_size_bytes(fp32_dir))
+    quantized_size_mb = _mb(_dir_size_bytes(int8_dir))
+    compression_ratio = (
+        float(original_size_mb) / float(quantized_size_mb)
+        if quantized_size_mb > 0
+        else 0.0
+    )
+
+    info = {
+        "original_model": checkpoint_dir_name,
+        "quantization_type": "dynamic_int8",
+        "original_size_mb": round(original_size_mb, 2),
+        "quantized_size_mb": round(quantized_size_mb, 2),
+        "compression_ratio": round(compression_ratio, 3),
+    }
+    info_path = os.path.join(int8_dir, "quantization_info.json")
+    with open(info_path, "w", encoding="utf-8") as f:
+        json.dump(info, f, indent=2)
+    return {"model_path": model_path, "info_path": info_path, "info": info}
+
+
+def _print_table(
+    fp32_size_mb: float,
+    int8_size_mb: float,
+    fp32_single_ms: float,
+    int8_single_ms: float,
+    fp32_batch16_ms: float,
+    int8_batch16_ms: float,
+    fp32_acc: float,
+    int8_acc: float,
+) -> None:
+    size_change_pct = 100.0 * (1.0 - (int8_size_mb / fp32_size_mb)) if fp32_size_mb > 0 else 0.0
+    single_speedup = (fp32_single_ms / int8_single_ms) if int8_single_ms > 0 else 0.0
+    batch_speedup = (fp32_batch16_ms / int8_batch16_ms) if int8_batch16_ms > 0 else 0.0
+    acc_delta_pp = (int8_acc - fp32_acc) * 100.0
+
+    def line(a: str, b: str, c: str, d: str) -> str:
+        return f"│ {a:<15} │ {b:<10} │ {c:<11} │ {d:<17} │"
+
+    print("┌─────────────────┬────────────┬─────────────┬───────────────────┐")
+    print(line("Metric", "FP32 Model", "INT8 Model", "Change"))
+    print("├─────────────────┼────────────┼─────────────┼───────────────────┤")
+    print(
+        line(
+            "Model size",
+            f"{fp32_size_mb:.1f} MB",
+            f"{int8_size_mb:.1f} MB",
+            f"-{size_change_pct:.1f}% smaller",
+        )
+    )
+    print(
+        line(
+            "Single-text ms",
+            f"{fp32_single_ms:.2f} ms",
+            f"{int8_single_ms:.2f} ms",
+            f"{single_speedup:.2f}x faster",
+        )
+    )
+    print(
+        line(
+            "Batch-16 ms",
+            f"{fp32_batch16_ms:.2f} ms",
+            f"{int8_batch16_ms:.2f} ms",
+            f"{batch_speedup:.2f}x faster",
+        )
+    )
+    print(
+        line(
+            "Test accuracy",
+            f"{fp32_acc * 100:.2f}%",
+            f"{int8_acc * 100:.2f}%",
+            f"{acc_delta_pp:+.2f} pp",
+        )
+    )
+    print("└─────────────────┴────────────┴─────────────┴───────────────────┘")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Dynamic INT8 quantization for transformer inference on CPU.")
+    parser.add_argument("--model", type=str, default="distilbert-base-uncased")
+    parser.add_argument("--benchmark-only", action="store_true")
+    args = parser.parse_args()
+
+    dir_name = _checkpoint_to_dir(args.model)
+    fp32_dir = os.path.join(CFG.models_dir, dir_name)
+    int8_dir = os.path.join(CFG.models_dir, f"{dir_name}_int8")
+
+    if not os.path.isdir(fp32_dir):
+        raise FileNotFoundError(
+            f"FP32 model directory not found: {fp32_dir}\n"
+            f"Expected a fine-tuned model saved via save_pretrained() under saved_models/."
+        )
+
+    print(f"[Quantize] Loading FP32 model from: {fp32_dir}")
+    model_fp32, tokenizer_fp32 = _load_fp32_model(fp32_dir)
+
+    print("[Quantize] Applying dynamic INT8 quantization (Linear layers)...")
+    model_int8 = _quantize_dynamic_int8(model_fp32)
+
+    if not args.benchmark_only:
+        saved = _save_quantized_model(model_int8, fp32_dir, int8_dir, checkpoint_dir_name=dir_name)
+        print(f"[Quantize] Saved INT8 model -> {saved['model_path']}")
+        print(f"[Quantize] Saved metadata  -> {saved['info_path']}")
+    else:
+        os.makedirs(int8_dir, exist_ok=True)
+
+    X_test, y_test = load_test_only()
+    rng = np.random.default_rng(CFG.seed)
+    sample_idx = rng.choice(len(X_test), size=min(100, len(X_test)), replace=False).tolist()
+    sample_texts = [X_test[i] for i in sample_idx]
+
+    print("[Benchmark] Measuring latency (median ms per text)...")
+    fp32_single = _benchmark_latency_ms(model_fp32, tokenizer_fp32, sample_texts, batch_size=1)
+    int8_single = _benchmark_latency_ms(model_int8, tokenizer_fp32, sample_texts, batch_size=1)
+    fp32_b16 = _benchmark_latency_ms(model_fp32, tokenizer_fp32, sample_texts, batch_size=16)
+    int8_b16 = _benchmark_latency_ms(model_int8, tokenizer_fp32, sample_texts, batch_size=16)
+
+    print("[Eval] Computing test accuracy on 7,600 examples...")
+    fp32_acc = _accuracy(model_fp32, tokenizer_fp32, X_test, y_test, batch_size=32)
+    int8_acc = _accuracy(model_int8, tokenizer_fp32, X_test, y_test, batch_size=32)
+
+    fp32_size_mb = _mb(_dir_size_bytes(fp32_dir))
+    int8_size_mb = _mb(_dir_size_bytes(int8_dir))
+
+    _print_table(
+        fp32_size_mb=fp32_size_mb,
+        int8_size_mb=int8_size_mb,
+        fp32_single_ms=fp32_single,
+        int8_single_ms=int8_single,
+        fp32_batch16_ms=fp32_b16,
+        int8_batch16_ms=int8_b16,
+        fp32_acc=fp32_acc,
+        int8_acc=int8_acc,
+    )
+
+
+if __name__ == "__main__":
+    main()
+

requirements.txt

@@ -0,0 +1,43 @@
+# PyTorch — install separately for your platform:
+#   pip install torch torchvision torchaudio
+# (MPS support is built-in on Apple Silicon; no extra steps needed)
+
+# ── HuggingFace Stack ────────────────────────────────────────────────────────
+transformers>=4.41.2
+datasets>=2.19.2
+evaluate>=0.4.2
+accelerate>=0.30.1
+huggingface-hub>=0.23.2
+tokenizers>=0.19.1
+
+# ── Traditional ML ───────────────────────────────────────────────────────────
+scikit-learn>=1.5.0
+joblib>=1.4.2
+
+# ── Data & Numerics ──────────────────────────────────────────────────────────
+numpy>=1.26.4
+pandas>=2.2.2
+
+# ── Visualisation ────────────────────────────────────────────────────────────
+# NOTE: matplotlib==3.9.0 fails to build from source on Apple clang 17
+#       (freetype-2.6.1 zconf.h Byte type conflict). Use >=3.9.0 to get
+#       a pre-built wheel for Python 3.13 arm64.
+matplotlib>=3.9.0
+seaborn>=0.13.2
+
+# ── API Server ───────────────────────────────────────────────────────────────
+fastapi>=0.111.0
+uvicorn[standard]>=0.30.1
+pydantic>=2.7.1
+
+# ── Optimisation ─────────────────────────────────────────────────────────────
+scipy>=1.14.0
+optuna==3.6.1
+
+# ── Misc ─────────────────────────────────────────────────────────────────────
+tqdm>=4.66.4
+
+# ── Testing ──────────────────────────────────────────────────────────────────
+pytest==8.2.2
+pytest-cov==5.0.0
+httpx==0.27.0

saved_models/distilbert_base_uncased/config.json

@@ -0,0 +1,39 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForSequenceClassification"
+  ],
+  "attention_dropout": 0.1,
+  "bos_token_id": null,
+  "dim": 768,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "hidden_dim": 3072,
+  "id2label": {
+    "0": "World",
+    "1": "Sports",
+    "2": "Business",
+    "3": "Sci/Tech"
+  },
+  "initializer_range": 0.02,
+  "label2id": {
+    "Business": 2,
+    "Sci/Tech": 3,
+    "Sports": 1,
+    "World": 0
+  },
+  "max_position_embeddings": 512,
+  "model_type": "distilbert",
+  "n_heads": 12,
+  "n_layers": 6,
+  "pad_token_id": 0,
+  "qa_dropout": 0.1,
+  "seq_classif_dropout": 0.2,
+  "sinusoidal_pos_embds": false,
+  "tie_weights_": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.12.0",
+  "use_cache": false,
+  "vocab_size": 30522
+}

saved_models/distilbert_base_uncased/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1579bcd6f3bd4d449ff1819d57a848097d37ba21ff1d4750bddbd61295b497de
+size 267838720

saved_models/distilbert_base_uncased/tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "backend": "tokenizers",
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

saved_models/distilbert_base_uncased/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:890e66b5a5ba897ccba31dc331ded90ea3933d7f8a010d0bec5b79c36c87b23f
+size 5201

saved_models/distilbert_base_uncased_int8/config.json

@@ -0,0 +1,39 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForSequenceClassification"
+  ],
+  "attention_dropout": 0.1,
+  "bos_token_id": null,
+  "dim": 768,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "hidden_dim": 3072,
+  "id2label": {
+    "0": "World",
+    "1": "Sports",
+    "2": "Business",
+    "3": "Sci/Tech"
+  },
+  "initializer_range": 0.02,
+  "label2id": {
+    "Business": 2,
+    "Sci/Tech": 3,
+    "Sports": 1,
+    "World": 0
+  },
+  "max_position_embeddings": 512,
+  "model_type": "distilbert",
+  "n_heads": 12,
+  "n_layers": 6,
+  "pad_token_id": 0,
+  "qa_dropout": 0.1,
+  "seq_classif_dropout": 0.2,
+  "sinusoidal_pos_embds": false,
+  "tie_weights_": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.12.0",
+  "use_cache": false,
+  "vocab_size": 30522
+}

saved_models/distilbert_base_uncased_int8/model_int8.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e77e69d402e37565c1ac4657663f065e211fcfde995e63f11cf1a0fcb055d9a3
+size 138718266

saved_models/distilbert_base_uncased_int8/quantization_info.json

@@ -0,0 +1,7 @@
+{
+  "original_model": "distilbert_base_uncased",
+  "quantization_type": "dynamic_int8",
+  "original_size_mb": 256.12,
+  "quantized_size_mb": 132.97,
+  "compression_ratio": 1.926
+}

saved_models/distilbert_base_uncased_int8/tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "backend": "tokenizers",
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

saved_models/traditional_lr.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76900d222138d705ba86d8337d9341179aa389064c95fd88651e3af6cc24128b
+size 4283180

saved_models/traditional_lr_optimized.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f343f5337ee10b32d436b598dbfbda2664d72d2e03495e28431cc458b19fc8f0
+size 4137004

saved_models/traditional_svm.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6254d1792546d51cc880d50a6a6e081bdb1c9830eed7073184d8df790fb25f66
+size 4283064

saved_models/traditional_svm_optimized.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50c72cd8331e30ba51fa4883dbf232cdf912cbb52799c4dffe60f84e898c585b
+size 5418824

tests/__init__.py

@@ -0,0 +1 @@
+

tests/conftest.py

@@ -0,0 +1,113 @@
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+import joblib
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.pipeline import Pipeline
+from sklearn.svm import LinearSVC
+
+import api
+import database
+from config import CFG
+
+
+@pytest.fixture
+def sample_texts():
+    return [
+        "UN peace talks resume as leaders meet to discuss ceasefire plans.",
+        "Local team wins championship after a dramatic overtime goal.",
+        "Stocks rally as the central bank signals a pause in rate hikes.",
+        "New smartphone chip boosts performance while reducing power use.",
+        "World leaders condemn attacks and call for immediate humanitarian aid.",
+        "Star striker sidelined with injury ahead of weekend match.",
+        "Company reports record quarterly earnings despite weak consumer demand.",
+        "Scientists discover a new exoplanet that may support liquid water.",
+        "Oil prices rise on supply concerns and geopolitical tensions.",
+        "Tech firm faces scrutiny after data breach exposes user accounts.",
+    ]
+
+
+@pytest.fixture
+def sample_labels():
+    return [0, 1, 2, 3, 0, 1, 2, 3, 2, 3]
+
+
+@pytest.fixture
+def mock_pipeline(sample_texts, sample_labels):
+    tfidf = TfidfVectorizer(
+        max_features=200,
+        ngram_range=(1, 1),
+        min_df=1,
+        sublinear_tf=False,
+    )
+    clf = LogisticRegression(
+        max_iter=200,
+        solver="lbfgs",
+        random_state=CFG.seed,
+    )
+    pipe = Pipeline([("tfidf", tfidf), ("lr", clf)])
+    pipe.fit(sample_texts, sample_labels)
+    return pipe
+
+
+@pytest.fixture
+def test_db_path(tmp_path):
+    return tmp_path / "test_requests.db"
+
+
+@pytest.fixture
+def api_client(monkeypatch, tmp_path, test_db_path):
+    models_dir = tmp_path / "models"
+    models_dir.mkdir(parents=True, exist_ok=True)
+    monkeypatch.setattr(CFG, "models_dir", str(models_dir))
+    monkeypatch.setattr(database, "_default_db_path", lambda: str(test_db_path))
+    database.init_db(db_path=str(test_db_path))
+
+    api._registry.clear()
+    test_app = FastAPI()
+    test_app.include_router(api.app.router)
+    return TestClient(test_app)
+
+
+@pytest.fixture
+def api_client_with_models(api_client, mock_pipeline, tmp_path, monkeypatch):
+    models_dir = tmp_path / "models"
+    monkeypatch.setattr(CFG, "models_dir", str(models_dir))
+    joblib.dump(mock_pipeline, models_dir / "traditional_lr.joblib")
+    svm = Pipeline(
+        [
+            ("tfidf", TfidfVectorizer(max_features=200, ngram_range=(1, 1), min_df=1)),
+            ("svm", LinearSVC(random_state=CFG.seed, max_iter=1000)),
+        ]
+    )
+    svm.fit(
+        [
+            "UN talks continue amid international pressure",
+            "Team wins match after extra time",
+            "Shares climb after earnings beat expectations",
+            "New processor improves phone battery life",
+            "Markets react to inflation report and central bank comments",
+            "Scientists unveil new telescope instrument",
+            "Player scores hat-trick in league game",
+            "Company announces merger in tech sector",
+        ],
+        [0, 1, 2, 3, 2, 3, 1, 2],
+    )
+    joblib.dump(svm, models_dir / "traditional_svm.joblib")
+    return api_client
+
+
+@pytest.fixture
+def api_client_with_startup(monkeypatch, tmp_path):
+    test_db_path = tmp_path / "startup_requests.db"
+    monkeypatch.setattr(database, "_default_db_path", lambda: str(test_db_path))
+    monkeypatch.setattr(CFG, "models_dir", str(tmp_path / "models"))
+    api._registry.clear()
+
+    def _fake_load_model(model_name: str):
+        raise FileNotFoundError("default model not available")
+
+    monkeypatch.setattr(api, "_load_model", _fake_load_model)
+    with TestClient(api.app) as client:
+        yield client

tests/test_api.py

@@ -0,0 +1,176 @@
+def test_health_returns_ok(api_client):
+    resp = api_client.get("/health")
+    assert resp.status_code == 200
+    assert resp.json()["status"] == "ok"
+
+
+def test_labels_returns_four_classes(api_client):
+    resp = api_client.get("/labels")
+    assert resp.status_code == 200
+    assert len(resp.json()["labels"]) == 4
+
+
+def test_predict_missing_model_returns_404(api_client):
+    resp = api_client.post(
+        "/predict", json={"text": "Some news article", "model_name": "nonexistent_model"}
+    )
+    assert resp.status_code == 404
+
+
+def test_predict_empty_text_returns_422(api_client):
+    resp = api_client.post("/predict", json={"text": "", "model_name": "lr"})
+    assert resp.status_code == 422
+
+
+def test_batch_predict_too_many_texts_returns_422(api_client):
+    resp = api_client.post(
+        "/batch_predict", json={"texts": ["text"] * 257, "model_name": "lr"}
+    )
+    assert resp.status_code == 422
+
+
+def test_analytics_summary_returns_valid_json(api_client):
+    resp = api_client.get("/analytics/summary")
+    assert resp.status_code == 200
+    assert "total_requests" in resp.json()
+
+
+def test_models_endpoint_lists_sklearn_model(api_client_with_models):
+    resp = api_client_with_models.get("/models")
+    assert resp.status_code == 200
+    payload = resp.json()
+    names = {m["name"] for m in payload["models"]}
+    assert "lr" in names
+
+
+def test_predict_lr_success_logs_request(api_client_with_models, test_db_path):
+    import database
+
+    resp = api_client_with_models.post(
+        "/predict", json={"text": "Fed raises interest rates by 50 bps", "model_name": "lr"}
+    )
+    assert resp.status_code == 200
+    body = resp.json()
+    assert "request_id" in body
+    assert "is_low_confidence" in body
+    assert "latency_ms" in body
+
+    history = database.get_request_history(db_path=str(test_db_path), limit=10)
+    assert len(history) == 1
+    assert history[0]["request_id"] == body["request_id"]
+
+
+def test_predict_lr_missing_joblib_returns_404(api_client):
+    resp = api_client.post(
+        "/predict", json={"text": "Fed raises rates", "model_name": "lr"}
+    )
+    assert resp.status_code == 404
+
+
+def test_batch_predict_lr_logs_each_item(api_client_with_models, test_db_path):
+    import database
+
+    texts = ["Apple unveils new AI chip", "Team wins the final match in overtime"]
+    resp = api_client_with_models.post(
+        "/batch_predict", json={"texts": texts, "model_name": "lr"}
+    )
+    assert resp.status_code == 200
+    body = resp.json()
+    assert body["count"] == 2
+    assert len(body["predictions"]) == 2
+
+    history = database.get_request_history(db_path=str(test_db_path), limit=10)
+    assert len(history) == 2
+    assert all(int(r["is_batch"]) == 1 for r in history)
+
+
+def test_predict_svm_success_uses_decision_function(api_client_with_models):
+    resp = api_client_with_models.post(
+        "/predict", json={"text": "Championship match ends in overtime", "model_name": "svm"}
+    )
+    assert resp.status_code == 200
+    body = resp.json()
+    assert body["probabilities"] is None
+    assert 0.0 <= float(body.get("latency_ms", 0.0))
+
+
+def test_batch_predict_svm_success(api_client_with_models):
+    resp = api_client_with_models.post(
+        "/batch_predict",
+        json={"texts": ["Markets rise on earnings", "New chip released"], "model_name": "svm"},
+    )
+    assert resp.status_code == 200
+    body = resp.json()
+    assert body["count"] == 2
+
+
+def test_predict_lr_cached_path(api_client_with_models, test_db_path):
+    import database
+
+    r1 = api_client_with_models.post(
+        "/predict", json={"text": "Stocks fall on inflation data", "model_name": "lr"}
+    )
+    r2 = api_client_with_models.post(
+        "/predict", json={"text": "Stocks fall on inflation data", "model_name": "lr"}
+    )
+    assert r1.status_code == 200
+    assert r2.status_code == 200
+    history = database.get_request_history(db_path=str(test_db_path), limit=10)
+    assert len(history) == 2
+
+
+def test_low_confidence_review_flow(api_client_with_models, test_db_path, monkeypatch):
+    import database
+    from config import CFG
+
+    monkeypatch.setattr(CFG, "low_confidence_threshold", 0.99)
+    resp = api_client_with_models.post(
+        "/predict", json={"text": "Mixed signals from markets after earnings report", "model_name": "lr"}
+    )
+    assert resp.status_code == 200
+    request_id = resp.json()["request_id"]
+
+    flags = api_client_with_models.get("/analytics/low_confidence").json()
+    assert any(f["request_id"] == request_id for f in flags)
+
+    patch_resp = api_client_with_models.patch(
+        f"/analytics/review/{request_id}", json={"note": "needs review"}
+    )
+    assert patch_resp.status_code == 200
+
+    reviewed_flags = api_client_with_models.get("/analytics/low_confidence?reviewed=true").json()
+    match = [f for f in reviewed_flags if f["request_id"] == request_id]
+    assert match
+    assert int(match[0]["reviewed"]) == 1
+    assert match[0]["review_note"] == "needs review"
+
+    history = database.get_request_history(db_path=str(test_db_path), limit=10)
+    assert history
+
+
+def test_export_flags_creates_files(api_client_with_models, tmp_path, monkeypatch):
+    import database
+    from config import CFG
+
+    monkeypatch.setattr(CFG, "low_confidence_threshold", 0.99)
+    api_client_with_models.post(
+        "/predict", json={"text": "Unclear headline with mixed topics", "model_name": "lr"}
+    )
+
+    out_dir = tmp_path / "low_confidence_review"
+
+    original = database.export_low_confidence_to_folder
+
+    def _export_override():
+        return original(output_dir=str(out_dir))
+
+    monkeypatch.setattr(database, "export_low_confidence_to_folder", _export_override)
+    resp = api_client_with_models.post("/analytics/export_flags")
+    assert resp.status_code == 200
+    payload = resp.json()
+    assert payload["exported"] >= 1
+
+
+def test_health_with_startup_runs(api_client_with_startup):
+    resp = api_client_with_startup.get("/health")
+    assert resp.status_code == 200

tests/test_config.py

@@ -0,0 +1,26 @@
+import os
+
+from config import CFG
+
+
+def test_device_string_valid():
+    assert CFG.device in {"mps", "cuda", "cpu"}
+
+
+def test_num_labels_matches_label_names():
+    assert CFG.num_labels == len(CFG.label_names)
+
+
+def test_dirs_created_on_init():
+    assert os.path.isdir(CFG.models_dir)
+    assert os.path.isdir(CFG.outputs_dir)
+    assert os.path.isdir(CFG.logs_dir)
+
+
+def test_label_names_correct():
+    assert CFG.label_names == ["World", "Sports", "Business", "Sci/Tech"]
+
+
+def test_max_length_positive():
+    assert CFG.max_length > 0
+

tests/test_data_loader.py

@@ -0,0 +1,39 @@
+import pytest
+
+from data_loader import get_raw_splits, load_ag_news, load_test_only
+
+
+pytestmark = pytest.mark.slow
+
+
+def test_load_test_only_returns_correct_types():
+    X_test, y_test = load_test_only()
+    assert isinstance(X_test, list)
+    assert isinstance(y_test, list)
+
+
+def test_load_test_only_length():
+    X_test, y_test = load_test_only()
+    assert len(X_test) == 7600
+    assert len(y_test) == 7600
+
+
+def test_load_test_only_labels_valid():
+    _, y_test = load_test_only()
+    assert all(y in {0, 1, 2, 3} for y in y_test)
+
+
+def test_load_ag_news_with_caps():
+    dataset = load_ag_news(max_train=100, max_eval=50, max_test=50)
+    assert len(dataset["train"]) == 100
+    assert len(dataset["validation"]) == 50
+    assert len(dataset["test"]) == 50
+
+
+def test_get_raw_splits_shapes():
+    dataset = load_ag_news(max_train=100, max_eval=50, max_test=50)
+    X_train, y_train, X_val, y_val, X_test, y_test = get_raw_splits(dataset)
+    assert len(X_train) == len(y_train) == 100
+    assert len(X_val) == len(y_val) == 50
+    assert len(X_test) == len(y_test) == 50
+

tests/test_database.py

@@ -0,0 +1,87 @@
+import sqlite3
+
+import database
+
+
+def _tables(db_path: str):
+    conn = sqlite3.connect(db_path)
+    try:
+        rows = conn.execute(
+            "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name;"
+        ).fetchall()
+        return [r[0] for r in rows]
+    finally:
+        conn.close()
+
+
+def test_init_db_creates_tables(test_db_path):
+    database.init_db(db_path=str(test_db_path))
+    tables = _tables(str(test_db_path))
+    assert "requests" in tables
+    assert "model_stats" in tables
+    assert "low_confidence_flags" in tables
+
+
+def test_log_request_creates_row(test_db_path):
+    database.init_db(db_path=str(test_db_path))
+    database.log_request(
+        db_path=str(test_db_path),
+        request_id="abc-123",
+        model_name="lr",
+        input_text="Fed raises rates",
+        predicted_label="Business",
+        predicted_label_id=2,
+        confidence=0.91,
+        latency_ms=12.4,
+        is_batch=False,
+    )
+    history = database.get_request_history(db_path=str(test_db_path), limit=10)
+    assert len(history) == 1
+    assert history[0]["predicted_label"] == "Business"
+
+
+def test_low_confidence_flag(test_db_path, monkeypatch):
+    monkeypatch.setattr(database.CFG, "low_confidence_threshold", 0.60)
+    database.init_db(db_path=str(test_db_path))
+    database.log_request(
+        db_path=str(test_db_path),
+        request_id="low-1",
+        model_name="lr",
+        input_text="Fed raises rates",
+        predicted_label="Business",
+        predicted_label_id=2,
+        confidence=0.45,
+        latency_ms=12.4,
+        is_batch=False,
+    )
+    flags = database.get_low_confidence_flags(db_path=str(test_db_path), reviewed=False)
+    assert len(flags) == 1
+    assert abs(float(flags[0]["confidence"]) - 0.45) < 1e-9
+
+
+def test_mark_reviewed(test_db_path, monkeypatch):
+    monkeypatch.setattr(database.CFG, "low_confidence_threshold", 0.60)
+    database.init_db(db_path=str(test_db_path))
+    database.log_request(
+        db_path=str(test_db_path),
+        request_id="low-2",
+        model_name="lr",
+        input_text="Fed raises rates",
+        predicted_label="Business",
+        predicted_label_id=2,
+        confidence=0.45,
+        latency_ms=12.4,
+        is_batch=False,
+    )
+    database.mark_reviewed("low-2", note="reviewed", db_path=str(test_db_path))
+    flags = database.get_low_confidence_flags(db_path=str(test_db_path), reviewed=True)
+    assert len(flags) == 1
+    assert int(flags[0]["reviewed"]) == 1
+    assert flags[0]["review_note"] == "reviewed"
+
+
+def test_get_summary_empty(test_db_path):
+    database.init_db(db_path=str(test_db_path))
+    summary = database.get_summary(db_path=str(test_db_path), days=7)
+    assert summary["total_requests"] == 0
+

tests/test_traditional_model.py

@@ -0,0 +1,42 @@
+import pytest
+from sklearn.pipeline import Pipeline
+
+from config import CFG
+from traditional_model import build_pipeline, load_model, save_model, train
+
+
+def test_build_pipeline_lr():
+    pipe = build_pipeline("lr")
+    assert isinstance(pipe, Pipeline)
+    assert len(pipe.steps) == 2
+
+
+def test_build_pipeline_svm():
+    pipe = build_pipeline("svm")
+    assert isinstance(pipe, Pipeline)
+    assert len(pipe.steps) == 2
+
+
+def test_build_pipeline_invalid_type():
+    with pytest.raises(ValueError):
+        build_pipeline("xgb")
+
+
+def test_train_returns_pipeline_and_float(sample_texts, sample_labels):
+    model, val_acc = train(sample_texts, sample_labels, sample_texts, sample_labels, "lr")
+    assert isinstance(model, Pipeline)
+    assert 0.0 <= float(val_acc) <= 1.0
+
+
+def test_pipeline_predict_returns_valid_labels(mock_pipeline, sample_texts):
+    preds = mock_pipeline.predict(sample_texts)
+    assert all(int(p) in {0, 1, 2, 3} for p in preds)
+
+
+def test_save_and_load_model(mock_pipeline, tmp_path, monkeypatch):
+    monkeypatch.setattr(CFG, "models_dir", str(tmp_path))
+    save_model(mock_pipeline, "test")
+    loaded = load_model("test")
+    out = loaded.predict(["test text"])
+    assert out is not None
+

tests/test_transformer_model.py

@@ -0,0 +1,33 @@
+import numpy as np
+import pytest
+
+from config import CFG
+from transformer_model import _checkpoint_to_dir, compute_metrics, load_model
+
+
+def test_checkpoint_to_dir_hyphen():
+    assert _checkpoint_to_dir("roberta-base") == "roberta_base"
+
+
+def test_checkpoint_to_dir_slash():
+    assert _checkpoint_to_dir("org/model") == "org_model"
+
+
+def test_compute_metrics_perfect():
+    eval_pred = (np.array([[10, 0, 0, 0], [0, 10, 0, 0]]), np.array([0, 1]))
+    result = compute_metrics(eval_pred)
+    assert result["accuracy"] == 1.0
+    assert result["f1_macro"] == 1.0
+
+
+def test_compute_metrics_all_wrong():
+    eval_pred = (np.array([[0, 10, 0, 0], [0, 0, 10, 0]]), np.array([0, 1]))
+    result = compute_metrics(eval_pred)
+    assert result["accuracy"] == 0.0
+
+
+def test_load_model_missing_raises(tmp_path, monkeypatch):
+    monkeypatch.setattr(CFG, "models_dir", str(tmp_path))
+    with pytest.raises(FileNotFoundError):
+        load_model("nonexistent-model")
+

traditional_model.py

@@ -0,0 +1,204 @@
+"""
+traditional_model.py
+────────────────────
+Approach A: TF-IDF vectorisation + scikit-learn classifiers.
+
+Provides two classifier options:
+    'lr'  → Logistic Regression  (supports probability scores)
+    'svm' → Linear SVM           (slightly faster, no probability output)
+"""
+import logging
+import os
+import time
+from typing import Dict, Literal, Tuple
+
+import joblib
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    confusion_matrix,
+)
+from sklearn.pipeline import Pipeline
+from sklearn.svm import LinearSVC
+
+from config import CFG
+
+logger = logging.getLogger(__name__)
+
+
+# ── Pipeline factory ──────────────────────────────────────────────────────────
+
+def build_pipeline(model_type: Literal["lr", "svm"] = "lr") -> Pipeline:
+    """
+    Build a TF-IDF → classifier sklearn Pipeline.
+
+    TF-IDF settings:
+    - max_features=60_000 : vocabulary cap; covers ~99 % of AG News tokens
+    - ngram_range=(1, 2)  : unigrams + bigrams capture short phrases
+    - sublinear_tf=True   : apply log(TF) to dampen very frequent terms
+    - min_df=2            : discard hapax legomena (appear only once)
+    """
+    tfidf = TfidfVectorizer(
+        max_features=60_000,
+        ngram_range=(1, 2),
+        sublinear_tf=True,
+        min_df=2,
+        strip_accents="unicode",
+        analyzer="word",
+        token_pattern=r"\w{1,}",
+    )
+
+    if model_type == "lr":
+        clf = LogisticRegression(
+            C=5.0,
+            max_iter=1_000,
+            solver="saga",
+            n_jobs=-1,
+            random_state=CFG.seed,
+        )
+    elif model_type == "svm":
+        clf = LinearSVC(
+            C=1.0,
+            max_iter=2_000,
+            random_state=CFG.seed,
+        )
+    else:
+        raise ValueError(
+            f"Unknown model_type '{model_type}'. Valid choices: 'lr', 'svm'."
+        )
+
+    pipeline = Pipeline([("tfidf", tfidf), (model_type, clf)])
+    logger.info(f"Pipeline: TF-IDF -> {clf.__class__.__name__}")
+    return pipeline
+
+
+# ── Training ──────────────────────────────────────────────────────────────────
+
+def train(
+    X_train, y_train,
+    X_val,   y_val,
+    model_type: str = "lr",
+) -> Tuple[Pipeline, float]:
+    """
+    Fit the pipeline and report validation accuracy.
+
+    Returns
+    -------
+    (fitted_pipeline, validation_accuracy)
+    """
+    pipeline = build_pipeline(model_type)
+
+    logger.info(f"Training {model_type.upper()} on {len(X_train):,} samples ...")
+    t0 = time.perf_counter()
+    pipeline.fit(X_train, y_train)
+    elapsed = time.perf_counter() - t0
+    logger.info(f"Training complete in {elapsed:.1f}s")
+
+    val_preds = pipeline.predict(X_val)
+    val_acc   = accuracy_score(y_val, val_preds)
+    logger.info(f"Validation accuracy: {val_acc * 100:.2f}%")
+
+    return pipeline, val_acc
+
+
+# ── Evaluation ────────────────────────────────────────────────────────────────
+
+def evaluate(
+    pipeline: Pipeline,
+    X_test,
+    y_test,
+    save_dir: str = None,
+) -> Dict:
+    """
+    Run the pipeline on the test set, print a full report and save the
+    confusion matrix.
+
+    Returns
+    -------
+    dict with keys: accuracy, report, confusion_matrix
+    """
+    preds  = pipeline.predict(X_test)
+    acc    = accuracy_score(y_test, preds)
+    cm     = confusion_matrix(y_test, preds)
+    report = classification_report(
+        y_test, preds,
+        target_names=CFG.label_names,
+        digits=4,
+    )
+
+    print("\n" + "=" * 60)
+    print("  TRADITIONAL MODEL  --  TEST SET RESULTS")
+    print("=" * 60)
+    print(f"  Accuracy : {acc * 100:.2f}%\n")
+    print(report)
+
+    _plot_confusion_matrix(
+        cm,
+        title="Traditional Model -- Confusion Matrix",
+        save_dir=save_dir,
+    )
+
+    return {"accuracy": acc, "report": report, "confusion_matrix": cm}
+
+
+# ── Persistence ───────────────────────────────────────────────────────────────
+
+def save_model(pipeline: Pipeline, name: str = "lr") -> str:
+    """Serialise the pipeline with joblib."""
+    path = os.path.join(CFG.models_dir, f"traditional_{name}.joblib")
+    joblib.dump(pipeline, path)
+    logger.info(f"Model saved  ->  {path}")
+    return path
+
+
+def load_model(name: str = "lr") -> Pipeline:
+    """Deserialise a saved pipeline."""
+    path = os.path.join(CFG.models_dir, f"traditional_{name}.joblib")
+    if not os.path.exists(path):
+        raise FileNotFoundError(
+            f"No saved model at '{path}'. "
+            f"Run: python train_traditional.py --model {name}"
+        )
+    pipeline = joblib.load(path)
+    logger.info(f"Model loaded <-  {path}")
+    return pipeline
+
+
+# ── Internal helpers ──────────────────────────────────────────────────────────
+
+def _plot_confusion_matrix(
+    cm: np.ndarray,
+    title: str,
+    save_dir: str = None,
+) -> None:
+    fig, ax = plt.subplots(figsize=(7, 6))
+    sns.heatmap(
+        cm,
+        annot=True,
+        fmt="d",
+        cmap="Blues",
+        xticklabels=CFG.label_names,
+        yticklabels=CFG.label_names,
+        linewidths=0.5,
+        ax=ax,
+    )
+    ax.set_xlabel("Predicted Label", fontsize=11)
+    ax.set_ylabel("True Label", fontsize=11)
+    ax.set_title(title, fontsize=13, fontweight="bold")
+    plt.tight_layout()
+
+    if save_dir:
+        os.makedirs(save_dir, exist_ok=True)
+        fig_path = os.path.join(save_dir, "confusion_matrix.png")
+        plt.savefig(fig_path, dpi=150)
+        logger.info(f"Confusion matrix ->  {fig_path}")
+
+    plt.show()
+    plt.close(fig)

train_multi.py

@@ -0,0 +1,118 @@
+"""
+train_multi.py
+──────────────
+Train multiple transformer architectures back-to-back.
+Ideal for an overnight run on Mac M4 — each model saves independently.
+Results are compared at the end in a clean table.
+
+Usage
+─────
+    # Train DistilBERT + BERT + RoBERTa (default)
+    python train_multi.py
+
+    # Train only DistilBERT and RoBERTa
+    python train_multi.py --models distilbert-base-uncased roberta-base
+
+    # Single model
+    python train_multi.py --models roberta-base
+"""
+import argparse
+import logging
+import time
+
+from config import CFG
+from data_loader import load_ag_news, get_tokenizer, tokenise_dataset
+import transformer_model as trm
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s  %(levelname)-8s  %(message)s",
+    datefmt="%H:%M:%S",
+)
+
+DEFAULT_MODELS = [
+    "distilbert-base-uncased",
+    "bert-base-uncased",
+    "roberta-base",
+]
+
+
+def train_single_architecture(checkpoint: str) -> dict:
+    """
+    End-to-end train and evaluate one transformer checkpoint.
+    Mutates CFG.model_checkpoint for the duration of the run.
+    """
+    print(f"\n{'━' * 60}")
+    print(f"  Model: {checkpoint}")
+    print(f"{'━' * 60}")
+
+    CFG.model_checkpoint = checkpoint   # Override for this run
+
+    dataset   = load_ag_news()          # Full 120K (no cap in updated config)
+    tokenizer = get_tokenizer()
+    tokenised = tokenise_dataset(dataset, tokenizer)
+
+    t0      = time.perf_counter()
+    trainer = trm.train(tokenised, tokenizer, checkpoint=checkpoint)
+    elapsed = time.perf_counter() - t0
+
+    save_dir = f"outputs/{trm._checkpoint_to_dir(checkpoint)}"
+    results  = trm.evaluate(trainer, tokenised,
+                            checkpoint=checkpoint, save_dir=save_dir)
+    trm.save_model(trainer, tokenizer, checkpoint=checkpoint)
+
+    h, rem = divmod(int(elapsed), 3600)
+    m, s   = divmod(rem, 60)
+
+    return {
+        "checkpoint": checkpoint,
+        "accuracy":   results["accuracy"],
+        "f1_macro":   results["metrics"].get("test_f1_macro", 0.0),
+        "time":       f"{h}h {m}m {s}s",
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Train multiple transformer architectures sequentially"
+    )
+    parser.add_argument(
+        "--models", nargs="+", default=DEFAULT_MODELS,
+        help="Space-separated list of HuggingFace checkpoint names",
+    )
+    args = parser.parse_args()
+
+    device_label = "MPS (Apple Metal)" if CFG.device == "mps" else CFG.device.upper()
+    print(f"\n  Multi-Architecture Training Session")
+    print(f"  Device  : {device_label}")
+    print(f"  Models  : {', '.join(args.models)}")
+    print(f"  Dataset : AG News — full 120,000 training examples\n")
+
+    all_results = []
+    session_t0  = time.perf_counter()
+
+    for checkpoint in args.models:
+        result = train_single_architecture(checkpoint)
+        all_results.append(result)
+        print(f"\n  ✓  {checkpoint}  —  acc={result['accuracy']*100:.2f}%  time={result['time']}\n")
+
+    session_elapsed = time.perf_counter() - session_t0
+    h, rem = divmod(int(session_elapsed), 3600)
+    m, s   = divmod(rem, 60)
+
+    print(f"\n{'═' * 66}")
+    print(f"  {'Architecture':<28}  {'Accuracy':>10}  {'F1-Macro':>10}  {'Time':>10}")
+    print(f"{'─' * 66}")
+    for r in sorted(all_results, key=lambda x: x["accuracy"], reverse=True):
+        name = r["checkpoint"].split("/")[-1]
+        star = "  ◀ best" if r == max(all_results, key=lambda x: x["accuracy"]) else ""
+        print(
+            f"  {name:<28}  {r['accuracy']*100:>9.2f}%  "
+            f"{r['f1_macro']:>10.4f}  {r['time']:>10}{star}"
+        )
+    print(f"{'═' * 66}")
+    print(f"\n  Total session time: {h}h {m}m {s}s\n")
+
+
+if __name__ == "__main__":
+    main()

train_traditional.py

@@ -0,0 +1,81 @@
+"""
+train_traditional.py
+────────────────────
+Entry-point: trains TF-IDF + Logistic Regression or Linear SVM.
+
+Usage
+─────
+    python train_traditional.py                 # Logistic Regression (default)
+    python train_traditional.py --model svm     # Linear SVM
+    python train_traditional.py --full          # Use all 120 K training samples
+    python train_traditional.py --model svm --full
+"""
+import argparse
+import logging
+
+from config import CFG
+from data_loader import load_ag_news, get_raw_splits
+import traditional_model as tm
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s  %(levelname)-8s  %(message)s",
+    datefmt="%H:%M:%S",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Train traditional ML document classifier (TF-IDF + LR or SVM)"
+    )
+    p.add_argument(
+        "--model",
+        default="lr",
+        choices=["lr", "svm"],
+        help="'lr' = Logistic Regression  |  'svm' = Linear SVM  (default: lr)",
+    )
+    p.add_argument(
+        "--full",
+        action="store_true",
+        help="Disable sample cap; use all 120 K training examples (slower)",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    max_train = None if args.full else CFG.max_train_samples
+    max_eval  = None if args.full else CFG.max_eval_samples
+    sample_label = "Full dataset" if (args.full or max_train is None) else f"{max_train:,} samples (subset)"
+
+    print(f"\n{'-' * 60}")
+    print(f"  Document Classifier  --  Traditional ML Training")
+    print(f"  Model   : {args.model.upper()}")
+    print(f"  Samples : {sample_label}")
+    print(f"{'-' * 60}\n")
+
+    # Load data
+    dataset = load_ag_news(max_train=max_train, max_eval=max_eval, max_test=None)
+    X_train, y_train, X_val, y_val, X_test, y_test = get_raw_splits(dataset)
+
+    # Train
+    pipeline, val_acc = tm.train(
+        X_train, y_train,
+        X_val,   y_val,
+        model_type=args.model,
+    )
+
+    # Evaluate on test set
+    save_dir = f"outputs/{args.model}"
+    results  = tm.evaluate(pipeline, X_test, y_test, save_dir=save_dir)
+
+    # Save the model
+    tm.save_model(pipeline, name=args.model)
+
+    print(f"\n  Final test accuracy : {results['accuracy'] * 100:.2f}%")
+    print(f"  Model saved to      : saved_models/traditional_{args.model}.joblib\n")
+
+
+if __name__ == "__main__":
+    main()

train_transformer.py

@@ -0,0 +1,72 @@
+"""
+train_transformer.py
+────────────────────
+Entry-point: fine-tunes DistilBERT on AG News.
+
+Usage
+─────
+    python train_transformer.py           # 20 K subset  ≈  1.5–2.5 hrs on i3
+    python train_transformer.py --full    # 120 K full   ≈  8–10 hrs on i3
+
+Tip: Start with the subset to verify everything works, then run --full overnight.
+"""
+import argparse
+import logging
+
+from config import CFG
+from data_loader import load_ag_news, get_tokenizer, tokenise_dataset
+import transformer_model as trm
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s  %(levelname)-8s  %(message)s",
+    datefmt="%H:%M:%S",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Fine-tune DistilBERT document classifier"
+    )
+    p.add_argument(
+        "--full",
+        action="store_true",
+        help="Use the full 120 K training set instead of the 20 K subset",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    max_train = None if args.full else CFG.max_train_samples
+    max_eval  = None if args.full else CFG.max_eval_samples
+    sample_label = "Full dataset" if args.full else f"{max_train:,} samples (subset)"
+
+    print(f"\n{'─' * 60}")
+    print(f"  Document Classifier  —  DistilBERT Fine-Tuning")
+    print(f"  Base model : {CFG.model_checkpoint}")
+    print(f"  Samples    : {sample_label}")
+    print(f"{'─' * 60}\n")
+
+    # Load and tokenise dataset
+    dataset   = load_ag_news(max_train=max_train, max_eval=max_eval, max_test=None)
+    tokenizer = get_tokenizer()
+    tokenised = tokenise_dataset(dataset, tokenizer)
+
+    # Train
+    trainer = trm.train(tokenised, tokenizer)
+
+    # Evaluate on test set
+    save_dir = "outputs/transformer"
+    results  = trm.evaluate(trainer, tokenised, save_dir=save_dir)
+
+    # Save best checkpoint + tokeniser
+    trm.save_model(trainer, tokenizer)
+
+    print(f"\n  Final test accuracy : {results['accuracy'] * 100:.2f}%")
+    print(f"  Model saved to      : saved_models/transformer/\n")
+
+
+if __name__ == "__main__":
+    main()

transformer_model.py

@@ -0,0 +1,350 @@
+"""
+transformer_model.py
+────────────────────
+Updated for Mac M4 / Apple Silicon MPS.
+
+Key changes vs Windows version:
+  - Removed use_cpu=True → Trainer auto-detects MPS on Mac
+  - Added label_smoothing_factor
+  - Model-specific output directories (supports multiple architectures)
+  - Gradient checkpointing toggle
+  - Cleaned up device handling for inference
+"""
+import logging
+import os
+import time
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    confusion_matrix,
+    f1_score,
+)
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    EarlyStoppingCallback,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainingArguments,
+)
+
+from config import CFG
+
+logger = logging.getLogger(__name__)
+
+
+# ── Helper ────────────────────────────────────────────────────────────────────
+
+def _checkpoint_to_dir(checkpoint: str) -> str:
+    """Convert a HuggingFace checkpoint name to a safe directory name.
+
+    Examples:
+        'roberta-base'            → 'roberta_base'
+        'distilbert-base-uncased' → 'distilbert_base_uncased'
+    """
+    return checkpoint.replace("/", "_").replace("-", "_")
+
+
+# ── Model factory ─────────────────────────────────────────────────────────────
+
+def build_model(checkpoint: str = None) -> AutoModelForSequenceClassification:
+    """Load a pre-trained encoder with a randomly-initialised classification head."""
+    if checkpoint is None:
+        checkpoint = CFG.model_checkpoint
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        checkpoint,
+        num_labels=CFG.num_labels,
+        id2label={i: n for i, n in enumerate(CFG.label_names)},
+        label2id={n: i for i, n in enumerate(CFG.label_names)},
+    )
+
+    if CFG.use_gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+        logger.info("Gradient checkpointing: ON")
+
+    total     = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    logger.info(f"Model: {checkpoint} | total={total:,}  trainable={trainable:,}")
+    return model
+
+
+# ── Training arguments ────────────────────────────────────────────────────────
+
+def get_training_args(checkpoint: str = None, output_dir: str = None) -> TrainingArguments:
+    """
+    Build MPS-safe TrainingArguments for the HuggingFace Trainer.
+
+    Critical Mac M4 notes
+    ─────────────────────
+    • Do NOT set use_cpu=True — the Trainer auto-detects MPS on Mac
+    • fp16=False          — MPS lacks full float16 operator coverage
+    • bf16=False          — Keep False for reliability (can try True on M2+)
+    • dataloader_pin_memory=False — pin_memory only benefits CUDA
+    • dataloader_num_workers=0   — HuggingFace torch datasets + multiprocessing
+                                   can be unstable on Mac; 0 is safest
+
+    Transformers 5.x deprecations handled here
+    ──────────────────────────────────────────
+    • warmup_ratio → warmup_steps  (computed manually below)
+    • logging_dir  → TENSORBOARD_LOGGING_DIR env var
+    """
+    if checkpoint is None:
+        checkpoint = CFG.model_checkpoint
+    if output_dir is None:
+        output_dir = os.path.join(CFG.outputs_dir, _checkpoint_to_dir(checkpoint))
+
+    # Compute warmup_steps from ratio (replaces deprecated warmup_ratio arg)
+    total_steps   = (108_000 // CFG.batch_size) * CFG.num_epochs // CFG.grad_accum_steps
+    warmup_steps  = int(total_steps * CFG.warmup_ratio)
+
+    # Set TensorBoard log dir via env var (replaces deprecated logging_dir arg)
+    os.environ["TENSORBOARD_LOGGING_DIR"] = CFG.logs_dir
+
+    return TrainingArguments(
+        output_dir=output_dir,
+
+        # ── Schedule ────────────────────────────────────────────────────────
+        num_train_epochs=CFG.num_epochs,
+        per_device_train_batch_size=CFG.batch_size,
+        per_device_eval_batch_size=CFG.batch_size * 2,
+        gradient_accumulation_steps=CFG.grad_accum_steps,
+
+        # ── Optimiser ──────────────────────────────────────���─────────────────
+        learning_rate=CFG.learning_rate,
+        weight_decay=CFG.weight_decay,
+        warmup_steps=warmup_steps,           # replaces deprecated warmup_ratio
+        lr_scheduler_type="cosine",
+        label_smoothing_factor=CFG.label_smoothing,
+
+        # ── Evaluation & checkpointing ───────────────────────────────────────
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        metric_for_best_model="accuracy",
+        greater_is_better=True,
+        save_total_limit=2,
+
+        # ── Logging ──────────────────────────────────────────────────────────
+        # logging_dir is deprecated in transformers 5.x; use TENSORBOARD_LOGGING_DIR env var instead
+        logging_steps=100,
+        report_to="none",
+
+        # ── MPS / Mac-specific ────────────────────────────────────────────────
+        # NOTE: No use_cpu=True here — MPS is auto-detected on Mac M4
+        fp16=False,
+        bf16=False,
+        dataloader_num_workers=CFG.num_workers,
+        dataloader_pin_memory=False,
+
+        # ── Reproducibility ──────────────────────────────────────────────────
+        seed=CFG.seed,
+        data_seed=CFG.seed,
+        push_to_hub=False,
+    )
+
+
+# ── Metrics ───────────────────────────────────────────────────────────────────
+
+def compute_metrics(eval_pred) -> Dict[str, float]:
+    """Called by Trainer after every validation epoch."""
+    logits, labels = eval_pred
+    preds = np.argmax(logits, axis=-1)
+    return {
+        "accuracy": float(accuracy_score(labels, preds)),
+        "f1_macro": float(f1_score(labels, preds, average="macro")),
+    }
+
+
+# ── Training pipeline ─────────────────────────────────────────────────────────
+
+def train(tokenised_dataset, tokenizer: PreTrainedTokenizerBase,
+          checkpoint: str = None) -> Trainer:
+    """
+    Fine-tune a transformer encoder and return the Trainer
+    with the best checkpoint loaded.
+    """
+    if checkpoint is None:
+        checkpoint = CFG.model_checkpoint
+
+    model         = build_model(checkpoint)
+    training_args = get_training_args(checkpoint)
+    data_collator = DataCollatorWithPadding(tokenizer, return_tensors="pt")
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenised_dataset["train"],
+        eval_dataset=tokenised_dataset["validation"],
+        processing_class=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
+    )
+
+    steps_per_epoch = len(tokenised_dataset["train"]) // CFG.batch_size
+    device_label    = "MPS (Metal)" if CFG.device == "mps" else CFG.device.upper()
+
+    logger.info("═" * 60)
+    logger.info(f"  Fine-Tuning: {checkpoint}")
+    logger.info(f"  Device     : {device_label}")
+    logger.info(f"  train      : {len(tokenised_dataset['train']):,}")
+    logger.info(f"  val        : {len(tokenised_dataset['validation']):,}")
+    logger.info(f"  epochs     : {CFG.num_epochs}")
+    logger.info(f"  batch      : {CFG.batch_size}")
+    logger.info(f"  steps/ep   : {steps_per_epoch:,}")
+    logger.info(f"  max_length : {CFG.max_length}")
+    logger.info("═" * 60)
+
+    t0 = time.perf_counter()
+    trainer.train()
+    elapsed = time.perf_counter() - t0
+
+    h, rem = divmod(int(elapsed), 3600)
+    m, s   = divmod(rem, 60)
+    logger.info(f"Training complete: {h}h {m}m {s}s")
+    return trainer
+
+
+# ── Evaluation ────────────────────────────────────────────────────────────────
+
+def evaluate(trainer: Trainer, tokenised_dataset,
+             checkpoint: str = None, save_dir: str = None) -> Dict:
+    """Run predictions on the test split and print full report."""
+    if checkpoint is None:
+        checkpoint = CFG.model_checkpoint
+
+    logger.info(f"Evaluating {checkpoint} on test set …")
+    predictions = trainer.predict(tokenised_dataset["test"])
+
+    preds  = np.argmax(predictions.predictions, axis=-1)
+    labels = predictions.label_ids
+
+    acc    = accuracy_score(labels, preds)
+    report = classification_report(labels, preds,
+                                   target_names=CFG.label_names, digits=4)
+    cm     = confusion_matrix(labels, preds)
+
+    print("\n" + "═" * 60)
+    print(f"  {checkpoint.upper()}  —  TEST SET RESULTS")
+    print("═" * 60)
+    print(f"  Accuracy : {acc * 100:.2f}%")
+    print(f"  Metrics  : {predictions.metrics}\n")
+    print(report)
+
+    _plot_cm(cm, f"{checkpoint} — Confusion Matrix",
+             save_dir=save_dir, cmap="Greens")
+
+    return {
+        "accuracy":         acc,
+        "report":           report,
+        "confusion_matrix": cm,
+        "metrics":          predictions.metrics,
+    }
+
+
+# ── Persistence ───────────────────────────────────────────────────────────────
+
+def save_model(trainer: Trainer, tokenizer: PreTrainedTokenizerBase,
+               checkpoint: str = None) -> str:
+    """Save best checkpoint + tokeniser to saved_models/<model_dir>/."""
+    if checkpoint is None:
+        checkpoint = CFG.model_checkpoint
+    path = os.path.join(CFG.models_dir, _checkpoint_to_dir(checkpoint))
+    trainer.save_model(path)
+    tokenizer.save_pretrained(path)
+    logger.info(f"Model saved  →  {path}")
+    return path
+
+
+def load_model(checkpoint: str = None) -> Tuple:
+    """
+    Load a saved fine-tuned model and its tokeniser.
+
+    Parameters
+    ----------
+    checkpoint : HuggingFace checkpoint name, e.g. 'roberta-base'.
+                 If None, uses CFG.model_checkpoint.
+
+    Returns
+    -------
+    (model, tokenizer) — model is in eval mode
+    """
+    if checkpoint is None:
+        checkpoint = CFG.model_checkpoint
+    path = os.path.join(CFG.models_dir, _checkpoint_to_dir(checkpoint))
+    if not os.path.isdir(path):
+        raise FileNotFoundError(
+            f"No saved model at '{path}'.\n"
+            f"Run: python train_transformer.py  (or  python train_multi.py)"
+        )
+    model     = AutoModelForSequenceClassification.from_pretrained(path)
+    tokenizer = AutoTokenizer.from_pretrained(path)
+    model.eval()
+    logger.info(f"Model loaded ←  {path}")
+    return model, tokenizer
+
+
+def load_quantized_model(checkpoint: str = "distilbert-base-uncased") -> Tuple:
+    """
+    Load the INT8 dynamically quantized version of a model.
+    Falls back to the FP32 model if the INT8 version is not found.
+
+    Returns
+    -------
+    (model, tokenizer, is_quantized)
+    """
+    dir_name = _checkpoint_to_dir(checkpoint)
+    int8_path = os.path.join(CFG.models_dir, f"{dir_name}_int8")
+    fp32_path = os.path.join(CFG.models_dir, dir_name)
+    model_file = os.path.join(int8_path, "model_int8.pt")
+
+    if os.path.exists(model_file):
+        # Apple Silicon/ARM requires the qengine (qnnpack) to be set before deserialising
+        try:
+            torch.backends.quantized.engine = "qnnpack"
+        except Exception:
+            pass
+
+        try:
+            model = torch.load(model_file, map_location="cpu", weights_only=False)
+        except TypeError:
+            model = torch.load(model_file, map_location="cpu")
+        tokenizer = AutoTokenizer.from_pretrained(int8_path)
+        model.eval()
+        logger.info(f"INT8 quantized model loaded ← {int8_path}")
+        return model, tokenizer, True
+
+    logger.warning(f"INT8 model not found at {int8_path}. Falling back to FP32.")
+    model, tokenizer = load_model(checkpoint)
+    return model, tokenizer, False
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _plot_cm(cm: np.ndarray, title: str,
+             save_dir: str = None, cmap: str = "Blues") -> None:
+    fig, ax = plt.subplots(figsize=(7, 6))
+    sns.heatmap(cm, annot=True, fmt="d", cmap=cmap,
+                xticklabels=CFG.label_names, yticklabels=CFG.label_names,
+                linewidths=0.5, ax=ax)
+    ax.set_xlabel("Predicted Label", fontsize=11)
+    ax.set_ylabel("True Label", fontsize=11)
+    ax.set_title(title, fontsize=13, fontweight="bold")
+    plt.tight_layout()
+    if save_dir:
+        os.makedirs(save_dir, exist_ok=True)
+        path = os.path.join(save_dir, "confusion_matrix.png")
+        plt.savefig(path, dpi=150)
+        logger.info(f"Confusion matrix →  {path}")
+    plt.show()
+    plt.close(fig)