Skip to main content

ML Classifier Service

Overview​

The ML classifier is a self-hosted FastAPI service running BAAI/bge-small-en-v1.5 with a 3-class classification head. It provides higher accuracy than the rule-tier V0 classifier (~75% vs ~60%).

Quick Start​

Docker Compose​

classifier:
  build:
    context: .
    dockerfile: deploy/classifier/Dockerfile
  ports:
    - "8082:8080"
  deploy:
    resources:
      limits: { memory: 2G, cpus: "2.0" }
  healthcheck:
    test: ["CMD", "curl", "-f", "http://localhost:8080/health"]
    interval: 30s
    timeout: 5s
    retries: 3

Manual Run​

pip install fastapi uvicorn transformers torchsentencepiece
python deploy/classifier/server.py --port 8082

API Endpoints​

POST /v1/classify​

Classify a prompt's complexity level.

Request:

{
  "text": "Debug this Python code:\n\ndef calculate_sum(numbers):\n    total = 0\n    for n in numbers:\n        total += n\n    return total",
  "org_id": "org-123"
}

Response:

{
  "class_name": "hard",
  "confidence": 0.87
}

GET /health​

Health check endpoint.

Response: {"status": "healthy", "model": "BAAI/bge-small-en-v1.5"}

Model Details​

  • Base model: BAAI/bge-small-en-v1.5 (137M parameters, 512 max tokens)
  • Head: 3-class linear classifier (easy/medium/hard)
  • Inference latency: ~20ms on CPU (single request)
  • Embedding dimension: 384 (used for semantic cache lookup)

Fine-tuning Pipeline​

Step 1: Build Dataset​

python scripts/build_dataset.py --output datasets/classifier/prompts_train.jsonl
# Generates 9k synthetic prompts (40% easy, 35% medium, 25% hard)

Step 2: Fine-tune​

python scripts/train_classifier.py \
    --train datasets/classifier/prompts_train.jsonl \
    --test datasets/classifier/prompts_test.jsonl \
    --output artifacts/classifier-v1.pt \
    --epochs 5 \
    --batch-size 8

Step 3: Export to ONNX​

pip install optimum onnxruntime
optimum export onnx \
    --model artifacts/classifier-v1.pt \
    --task text-classification \
    classifier.onnx

Step 4: Deploy​

Replace model.pt in the Docker image with the new ONNX export.

Performance Benchmarks​

MetricRule-tier V0ML V1
Accuracy (synthetic)~60%~75%+
Latency (p99)<1ms20-50ms
Confidence alignmentN/APlatt calibrated
Fallback chainN/ARule-tier V0

Circuit Breaker Integration​

The ML service is wrapped in a circuit breaker (see pkg/classifier/circuit_breaker.go). If the service fails 3 consecutive times, it enters Open state for 30 seconds — all requests fast-fail to rule-tier without HTTP calls. After 30s, one probe request tests recovery.

Scaling Notes​

  • CPU-only inference is sufficient for ~100 RPS
  • For >1000 RPS, consider ONNX Runtime with multiple threads or GPU acceleration
  • The embedding output (384-dim) is used for semantic cache lookup in pgvector