Hugging Face: DistilBERT
Fine-tune DistilBERT on an IMDB subset with the Trainer API. A TrainerCallback forwards the eval metrics; the datasets library loads the data locally and nothing about it leaves your machine.
Install
The framework, plus the K-Veritas CLI (see the Overview for install).
pip install transformers datasets accelerate torch
train.py
The complete script. The KVERITAS_ lines are the only additions to an ordinary training script; everything else is standard Hugging Face Transformers.
import numpy as np
import torch
from datasets import load_dataset
from transformers import (
AutoTokenizer, AutoModelForSequenceClassification,
TrainingArguments, Trainer, TrainerCallback,
)
SEED = 42
torch.manual_seed(SEED)
print(f"KVERITAS_INPUT src=seed:{SEED}", flush=True)
MODEL = "distilbert-base-uncased"
print("KVERITAS_PHASE name=data", flush=True)
raw = load_dataset("stanfordnlp/imdb")
train_ds = raw["train"].shuffle(seed=SEED).select(range(2000))
test_ds = raw["test"].shuffle(seed=SEED).select(range(1000))
tokenizer = AutoTokenizer.from_pretrained(MODEL)
def tok(batch):
return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=128)
train_ds = train_ds.map(tok, batched=True).rename_column("label", "labels")
test_ds = test_ds.map(tok, batched=True).rename_column("label", "labels")
model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)
n_params = sum(p.numel() for p in model.parameters())
print(f"KVERITAS_MODEL params={n_params} arch=distilbert precision=fp16", flush=True)
print(f"KVERITAS_WORKLOAD dataset_size={len(train_ds)} epochs=2 batch_size=16 seq_len=128", flush=True)
def compute_metrics(eval_pred):
logits, labels = eval_pred
preds = np.argmax(logits, axis=1)
return {"accuracy": float((preds == labels).mean())}
class KVeritas(TrainerCallback):
def on_evaluate(self, args, state, control, metrics=None, **kwargs):
for key, value in (metrics or {}).items():
if isinstance(value, (int, float)):
name = key.replace("eval_", "").replace("/", "_")
print(f"KVERITAS_METRIC name={name} value={float(value):.4f} step={int(state.epoch or 0)}", flush=True)
args = TrainingArguments(
output_dir="./_out",
per_device_train_batch_size=16,
per_device_eval_batch_size=32,
num_train_epochs=2,
eval_strategy="epoch",
logging_strategy="no",
save_strategy="no",
report_to="none",
fp16=torch.cuda.is_available(),
seed=SEED,
)
trainer = Trainer(
model=model, args=args,
train_dataset=train_ds, eval_dataset=test_ds,
processing_class=tokenizer, compute_metrics=compute_metrics,
callbacks=[KVeritas()],
)
print("KVERITAS_PHASE name=train", flush=True)
trainer.train()
print("KVERITAS_PHASE name=evaluate", flush=True)
final = trainer.evaluate()
acc = float(final["eval_accuracy"])
print(f"KVERITAS_METRIC name=test_accuracy value={acc:.4f}", flush=True)
print(f"KVERITAS_CLAIM metric=accuracy value={acc:.4f}", flush=True)
print(f"final accuracy={acc:.4f}", flush=True)Run and seal
Wrap the script, seal a signed report, and verify it offline.
kveritas init kveritas run -- python train.py kveritas seal --output report.pdf kveritas verify report.pdf
What it produced here
Two epochs fine-tuning DistilBERT on a 2,000-example IMDB subset reached about 83% accuracy; the callback sealed the eval metrics each epoch.