polyglot-tutor / training /README.md
Arthur_Diaz
feat(ml): ONNX int8 export and torch-free CPU inference service (#4)
933025e unverified
|
Raw
History Blame Contribute Delete
2.56 kB

training/

GPU-side code for the M1 CEFR classifier. Runs on the local box (WSL2, RTX 3070) with the train dependency group and never enters the runtime Docker image (ADR 0001). The record policies it applies — label mapping, chunking, document-level split, aggregation, metrics — live in src/tutor/ml/cefr/ so inference shares them byte-for-byte (no train/serve skew) and CI tests them without torch.

One-time setup

uv add --group train torch transformers accelerate datasets mlflow
uv run --group train python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"

Runs (ADR 0003 experiment arms)

# Validate a config without GPU (split summary + leakage check):
uv run --group train python training/train_cefr.py --config training/configs/en_only.toml --dry-run

# Arms 1-3:
uv run --group train python training/train_cefr.py --config training/configs/en_only.toml
uv run --group train python training/train_cefr.py --config training/configs/multilingual.toml
uv run --group train python training/train_cefr.py --config training/configs/en_truncated.toml

# Published baseline on the same test docs:
uv run --group train python training/eval_cefr.py \
  --model UniversalCEFR/xlm-roberta-base-cefr-all-classifier \
  --config training/configs/en_only.toml

Arm 5 (generalization audits) = copies of en_only.toml with exclude_corpora_from_train = ["cambridge_exams_en"] (and the elg twin).

Tracking

uv run --group train mlflow ui --backend-store-uri sqlite:///mlflow.db

Headline metrics: test_en_document_* (macro-F1, adjacent accuracy, QWK).

Export & deploy (M1)

uv add --group train onnx   # quantization dependency (export uses torch.onnx directly)

# Export + int8 + equivalence check + CPU bench (threads=2 ~ Space cpu-basic)
uv run --group train python training/export_onnx.py \
  --run-dir models/cefr/en_chunked_weighted --threads 2

# Parity: the deployed int8 service re-scored on the canonical test docs
uv run --group train python training/eval_service.py \
  --artifact models/cefr/en_chunked_weighted/onnx-int8 \
  --config training/configs/en_only.toml

# Publish the artifact (model card with cc-by-nc-sa-4.0 included)
uv run --group train python training/export_onnx.py \
  --run-dir models/cefr/en_chunked_weighted \
  --push <hf-username>/polyglot-tutor-cefr-onnx

The Space loads the artifact via CEFR_MODEL_ID (or CEFR_MODEL_PATH locally); the runtime depends only on onnxruntime + tokenizers — torch never ships.