mmBERT Multi-Dimensional Quality Classifier

Training pipeline for a multi-head quality classification model built on top of mmBERT. The model predicts 6 quality dimensions for a given text:

• clarity

• coherence

• depth

• grammar

• usefulness

• overall

Each dimension is a 3-class classification task: low, medium, high.

Overview

This repository provides:

• Training pipeline (train.py)

• Dataset processing and splitting (dataset.py)

• Model definition (model.py)

• Inference utilities (inference.py)

• SkyPilot launcher (sky_launch.py)

• Multi-language training configs (Tamil, Hindi, English)

The model uses a shared encoder with independent classification heads per dimension.

Architecture

• Base model: jhu-clsp/mmBERT-base

• Input: text

• Output: 6 independent classification heads (one per dimension)

Flow:

input text → tokenizer → mmBERT encoder → [CLS] embedding → dropout → 6 linear heads → predictions :contentReference[oaicite:0]{index=0}

Installation

Recommended: use uv

pip install uv --quiet
uv sync

Dependencies are defined in pyproject.toml

Authentication (hf_token.txt)

Pushing checkpoints and writing the leaderboard need a HuggingFace write token. Each user creates a gitignored hf_token.txt at the repo root containing just the raw token — train.py and the scripts/ tools read it automatically:

echo -n 'hf_xxxxxxxx' > hf_token.txt     # gitignored; one line, raw token

• Local runs / scripts: nothing else needed — the token is read from the file.

• Cluster runs: hf_token.txt is listed in each yaml's file_sync, so it's synced next to the code and read the same way on the cluster.

(--hf_token and the HF_TOKEN env var are still accepted as overrides, in that order of precedence.)

Data Source

There are two ways to feed data in, with the HuggingFace dataset as the default:

• HuggingFace dataset (default). Pre-made, pre-split parquet at samarthramesh/gaperon-distill (one config per language: english/hindi/tamil, splits train/validation/test). It is fetched once and cached as memory-mapped Arrow under --cache_dir, then reused across epochs — built for a cluster with no persistent storage and a slow link. Rows carry flat columns: text, doc_id, language, and the 6 label columns. Select it with --hf_repo/--lang/--cache_dir.

• Local .jsonl shards (offline override). Pass --data_dir to read .jsonl files from a directory and split them locally (see Data Splitting below). This bypasses the HF path entirely.

Local record format

Each .jsonl record should contain text and a judge object with 6 dimensions:

{
    "doc_id": "123",
    "text": "Sample text",
    "language": "en",
    "judge": {
        "clarity": "high",
        "coherence": "medium",
        "depth": "low",
        "grammar": "high",
        "usefulness": "medium",
        "overall": "high"
    }
}

Filtering

• Samples with any "unknown" label are dropped

• Only valid samples are used

Data Splitting

• Deterministic splits based on doc_id hash

• Default:

  • train: 80%

  • val: 10%

  • test: 10%

Generated files:

<data_dir>/splits/train.jsonl
<data_dir>/splits/val.jsonl
<data_dir>/splits/test.jsonl

Splits are reused unless --force_resplit is passed :contentReference[oaicite:3]{index=3}

Training

Basic Command (HF dataset, default)

python train.py \
  --hf_repo samarthramesh/gaperon-distill \
  --lang hindi \
  --cache_dir /scratch/hf_cache \
  --output_dir /path/to/models \
  --model_name my_model \
  --epochs 5

Offline override — read local .jsonl shards instead:

python train.py \
  --data_dir /path/to/shards \
  --output_dir /path/to/models \
  --model_name my_model \
  --epochs 5

Key Arguments

• --hf_repo / --lang / --cache_dir : HF dataset source (default path)

• --data_dir : local .jsonl shard directory (offline override)

• --output_dir : where checkpoints are saved

• --model_name : run identifier

• --epochs : number of epochs

• --batch_size : training batch size

• --lr : learning rate (default: 2e-5)

• --freeze_layers : freeze first N encoder layers

• --keep_last_n : checkpoint retention

• --limit : cap each split to N samples for a fast test run

• --push_to_hub (+ --hf_model_repo/--run_branch/--hf_token/--push_each_best) : push best/ to the Hub (see Model Versioning)

Full CLI defined in train.py

Checkpoints

Saved under:

<output_dir>/<model_name>/
  ├── best/
  ├── epoch-01/
  ├── epoch-02/
  └── ...

Contents

Each checkpoint includes:

• HuggingFace encoder weights

• tokenizer

• heads.pt (classification heads)

• training_info.json (metrics + config)

Best checkpoint is selected based on validation macro-F1 :contentReference[oaicite:5]{index=5}

Metrics

Per dimension:

• Macro F1

• Accuracy

Aggregate:

• mean_macro_F1 (primary selection metric)

Evaluation runs:

• every epoch (validation)

• once at end (test on best checkpoint) :contentReference[oaicite:6]{index=6}

Inference

Single Text

python inference.py \
  --checkpoint ./models/my_model/best \
  --text "Your input text"

Batch Inference

# HF dataset split (default)
python inference.py \
  --checkpoint ./models/my_model/best \
  --hf_repo samarthramesh/gaperon-distill --lang hindi \
  --split test \
  --output_path predictions.jsonl

# Local split file (offline override)
python inference.py \
  --checkpoint ./models/my_model/best \
  --data_dir /path/to/data \
  --split test \
  --output_path predictions.jsonl

Output includes predictions for all 6 dimensions.

Details in inference.py

Model Versioning (HF Hub)

Because the cluster has no persistent storage, trained checkpoints live on the HF Hub. Each language has its own model repo (samarthramesh/gaperon-mmbert-{english, hindi,tamil}); each run pushes its best/ checkpoint to its own branch, and the winning run is promoted to main.

Dev stage: while experimenting, push to the -dev repos (samarthramesh/gaperon-mmbert-{lang}-dev) via --hf_model_repo so the clean production repos stay reserved. The cluster yamls already point at the -dev repos.

1. Push a run to its branch — add --push_to_hub to training (token comes from hf_token.txt; see Authentication). By default the repo is gaperon-mmbert-<lang>; pass --hf_model_repo to override:

python train.py \
  --hf_repo samarthramesh/gaperon-distill --lang hindi \
  --output_dir /scratch/models --model_name v3_hindi_ordinal \
  --push_to_hub --hf_model_repo samarthramesh/gaperon-mmbert-hindi-dev
  # → gaperon-mmbert-hindi-dev @ runs/v3_hindi_ordinal

The branch carries the full checkpoint plus training_info.json (val + test metrics + hyperparameters), so every run is self-describing. By default the push happens once at the end of the run; add --push_each_best to also push on every new best (insurance against a job dying mid-run).

2. Compare runs — leaderboard.py reads every runs/* branch and ranks them by mean_macro_f1:

python scripts/leaderboard.py --repo samarthramesh/gaperon-mmbert-hindi
python scripts/leaderboard.py --repo samarthramesh/gaperon-mmbert-hindi \
  --sort_by test --write_readme     # also write the table into main README.md

3. Promote the winner — copy a run branch onto main:

python scripts/promote.py \
  --repo samarthramesh/gaperon-mmbert-hindi --branch runs/v3_hindi_ordinal

After promotion, from_pretrained("samarthramesh/gaperon-mmbert-hindi") returns that run's model.

Running on the Cluster (jclust)

The cluster has no persistent storage, so jobs pull data from HF (cached to /tmp) and push checkpoints back to the Hub. sky_launch.py syncs the files listed in the yaml's file_sync (including hf_token.txt) and merges setup into the run command.

Per-language configs are committed (train_{english,hindi,tamil}.yaml); each one streams from the HF dataset and pushes best/ to its -dev model repo. Create hf_token.txt (see Authentication) first, then launch:

python sky_launch.py jclust train_hindi.yaml

Dry-run to inspect the generated job without executing:

python sky_launch.py jclust train_hindi.yaml --dry-run

Auto teardown. Add --down to tear the cluster down once the run finishes (the checkpoint is already safe on the Hub, so the pod is disposable):

python sky_launch.py jclust train_hindi.yaml --down

--down only fires if this process stays alive until the job ends; the cluster also must be re-launched for the next run. For long jobs you may disconnect from, prefer a cluster-side timeout that survives disconnects:

sky autostop jclust --down -i 10     # auto-down 10 idle minutes after the job ends

Quick test run

train_test.yaml is a 1-epoch smoke test: it caps the data with --limit and pushes to the -dev repo, so the full path (load → train → eval → push) is exercised in seconds.

python sky_launch.py jclust train_test.yaml        # hindi, --limit 500, → …-hindi-dev

The same --limit N flag works locally for a fast end-to-end check:

python train.py \
  --hf_repo samarthramesh/gaperon-distill --lang english \
  --cache_dir ./.hf_cache --output_dir ./models \
  --model_name test_local_1ep --epochs 1 --batch_size 8 --limit 500 \
  --push_to_hub --hf_model_repo samarthramesh/gaperon-mmbert-english-dev

A --limit-ed run records "limit" in its training_info.json, so it's clearly marked as a test on the leaderboard.

Notes

• Dynamic padding is used per batch (no global padding)

• Classification heads are trained with higher LR than encoder

• Gradient clipping and LR scheduling are enabled

• Supports GPU and CPU execution

Minimal Workflow

1. Train from the HF dataset (--hf_repo/--lang) with --push_to_hub

2. Compare runs with scripts/leaderboard.py

3. Promote the winner to main with scripts/promote.py

4. Use the model via from_pretrained(<model_repo>) or run inference

(Offline variant: prepare .jsonl shards → train with --data_dir → use best/.)