Add training script

training_script.py

@@ -0,0 +1,214 @@
+# /// script
+# dependencies = [
+#   "torch",
+#   "transformers>=4.51.0",
+#   "datasets>=3.0.0",
+#   "accelerate>=1.0.0",
+#   "scikit-learn>=1.4.0",
+#   "trackio>=0.25.0",
+#   "huggingface_hub>=0.30.0",
+# ]
+# ///
+
+import os
+from collections import Counter
+
+import numpy as np
+import torch
+import trackio
+from datasets import load_dataset
+from huggingface_hub import HfApi
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+    set_seed,
+)
+
+DATASET_ID = "biglam/on_the_books"
+MODEL_ID = "distilbert-base-uncased"
+HUB_MODEL_ID = "evalstate/jim-crow-test2323"
+PROJECT = "jim-crow-law-classifier"
+RUN_NAME = "distilbert-on-the-books"
+MAX_LENGTH = 512
+SEED = 42
+
+set_seed(SEED)
+
+if not os.environ.get("HF_TOKEN"):
+    raise RuntimeError("HF_TOKEN is required so the trained model can be pushed to the Hub.")
+
+run = trackio.init(
+    project=PROJECT,
+    name=RUN_NAME,
+    config={
+        "dataset": DATASET_ID,
+        "base_model": MODEL_ID,
+        "hub_model_id": HUB_MODEL_ID,
+        "task": "binary sequence classification: Jim Crow law identification",
+        "max_length": MAX_LENGTH,
+        "seed": SEED,
+    },
+    private=False,
+    auto_log_gpu=True,
+)
+print(f"Trackio run: {run}")
+
+raw = load_dataset(DATASET_ID, split="train")
+label_names = raw.features["jim_crow"].names
+id2label = {i: name for i, name in enumerate(label_names)}
+label2id = {name: i for i, name in id2label.items()}
+print(raw)
+print("Label distribution:", Counter(raw["jim_crow"]))
+
+# Stratified split because the dataset has only one split and a modest class imbalance.
+splits = raw.train_test_split(test_size=0.2, seed=SEED, stratify_by_column="jim_crow")
+train_ds = splits["train"]
+eval_ds = splits["test"]
+
+trackio.log({
+    "data/train_examples": len(train_ds),
+    "data/eval_examples": len(eval_ds),
+    "data/train_jim_crow": Counter(train_ds["jim_crow"])[1],
+    "data/train_no_jim_crow": Counter(train_ds["jim_crow"])[0],
+})
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+def make_text(example):
+    chapter = example.get("chapter_text") or ""
+    section = example.get("section_text") or ""
+    meta = f"Source: {example.get('source','')}; Type: {example.get('type','')}; Chapter: {example.get('chapter_num','')}; Section: {example.get('section_num','')}"
+    return meta + "\n\nChapter text:\n" + chapter + "\n\nSection text:\n" + section
+
+def preprocess(batch):
+    texts = []
+    for i in range(len(batch["section_text"])):
+        ex = {k: batch[k][i] for k in batch.keys()}
+        texts.append(make_text(ex))
+    enc = tokenizer(texts, truncation=True, max_length=MAX_LENGTH)
+    enc["labels"] = batch["jim_crow"]
+    return enc
+
+remove_cols = raw.column_names
+train_tok = train_ds.map(preprocess, batched=True, remove_columns=remove_cols)
+eval_tok = eval_ds.map(preprocess, batched=True, remove_columns=remove_cols)
+
+counts = Counter(train_ds["jim_crow"])
+total = sum(counts.values())
+class_weights = torch.tensor([total / (2 * counts[i]) for i in range(len(label_names))], dtype=torch.float)
+print("Class weights:", class_weights.tolist())
+
+model = AutoModelForSequenceClassification.from_pretrained(
+    MODEL_ID,
+    num_labels=len(label_names),
+    id2label=id2label,
+    label2id=label2id,
+)
+
+class WeightedTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        labels = inputs.pop("labels")
+        outputs = model(**inputs)
+        weights = class_weights.to(outputs.logits.device)
+        loss_fct = torch.nn.CrossEntropyLoss(weight=weights)
+        loss = loss_fct(outputs.logits.view(-1, model.config.num_labels), labels.view(-1))
+        return (loss, outputs) if return_outputs else loss
+
+class TrackioCallback(TrainerCallback):
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs:
+            trackio.log({f"trainer/{k}": v for k, v in logs.items() if isinstance(v, (int, float))}, step=state.global_step)
+    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+        if metrics:
+            trackio.log({f"eval/{k}": v for k, v in metrics.items() if isinstance(v, (int, float))}, step=state.global_step)
+
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    preds = np.argmax(logits, axis=-1)
+    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average="binary", pos_label=1, zero_division=0)
+    macro_precision, macro_recall, macro_f1, _ = precision_recall_fscore_support(labels, preds, average="macro", zero_division=0)
+    acc = accuracy_score(labels, preds)
+    cm = confusion_matrix(labels, preds, labels=[0, 1])
+    return {
+        "accuracy": acc,
+        "precision": precision,
+        "recall": recall,
+        "f1": f1,
+        "macro_precision": macro_precision,
+        "macro_recall": macro_recall,
+        "macro_f1": macro_f1,
+        "tn": int(cm[0, 0]),
+        "fp": int(cm[0, 1]),
+        "fn": int(cm[1, 0]),
+        "tp": int(cm[1, 1]),
+    }
+
+args = TrainingArguments(
+    output_dir="jim-crow-test2323",
+    learning_rate=2e-5,
+    per_device_train_batch_size=16,
+    per_device_eval_batch_size=32,
+    gradient_accumulation_steps=1,
+    num_train_epochs=5,
+    weight_decay=0.01,
+    warmup_ratio=0.1,
+    lr_scheduler_type="linear",
+    eval_strategy="epoch",
+    save_strategy="epoch",
+    logging_steps=10,
+    load_best_model_at_end=True,
+    metric_for_best_model="f1",
+    greater_is_better=True,
+    save_total_limit=2,
+    fp16=torch.cuda.is_available(),
+    push_to_hub=True,
+    hub_model_id=HUB_MODEL_ID,
+    hub_private_repo=False,
+    report_to=[],
+    run_name=RUN_NAME,
+    seed=SEED,
+)
+
+trainer = WeightedTrainer(
+    model=model,
+    args=args,
+    train_dataset=train_tok,
+    eval_dataset=eval_tok,
+    processing_class=tokenizer,
+    data_collator=DataCollatorWithPadding(tokenizer),
+    compute_metrics=compute_metrics,
+    callbacks=[TrackioCallback()],
+)
+
+trainer.train()
+metrics = trainer.evaluate()
+print("Final eval metrics:", metrics)
+trackio.log({f"final/{k}": v for k, v in metrics.items() if isinstance(v, (int, float))})
+
+# Ensure useful metadata and a model card are present on the final Hub repo.
+trainer.save_model()
+tokenizer.save_pretrained(args.output_dir)
+trainer.create_model_card(
+    model_name="Jim Crow law classifier",
+    dataset_tags=DATASET_ID,
+    finetuned_from=MODEL_ID,
+    tasks="text-classification",
+    language="en",
+    tags=["legal", "history", "jim-crow", "sequence-classification", "distilbert"],
+)
+trainer.push_to_hub(commit_message="Fine-tune DistilBERT to identify Jim Crow laws")
+
+api = HfApi(token=os.environ["HF_TOKEN"])
+api.upload_file(
+    path_or_fileobj=__file__,
+    path_in_repo="training_script.py",
+    repo_id=HUB_MODEL_ID,
+    repo_type="model",
+    commit_message="Add training script",
+)
+print(f"Pushed trained model to https://huggingface.co/{HUB_MODEL_ID}")