Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,296 @@
+"""Full fine-tune of Whisper-small on the oddadmix multi-dialect Arabic set.
+
+  * Base:        openai/whisper-small (244M), full fine-tune (no LoRA).
+  * Data:        oddadmix/dialectal-arabic-lahgtna-v2-smaller-augmented
+                 (already 16 kHz mono; has an `augmentation` column).
+  * Targets:     cleaned with normalize.clean_text (tashkil + tags stripped).
+  * Features:    log-mel computed ON THE FLY in the collator, so we never
+                 materialize ~38 GB of cached features to disk.
+  * Precision:   bf16 on a 32 GB GPU.
+
+Auth: the dataset is private. Export a token first:
+    export HF_TOKEN=hf_xxx      # account with access to oddadmix/...
+Run:
+    uv run accelerate launch train.py            # or: uv run python train.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import torch
+from datasets import load_dataset
+import jiwer
+
+from transformers import (
+    WhisperProcessor,
+    WhisperForConditionalGeneration,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+)
+
+from normalize import clean_text
+
+MAX_AUDIO_SECONDS = 30.0   # Whisper encoder hard limit
+MIN_AUDIO_SECONDS = 0.5
+MAX_LABEL_TOKENS = 448     # Whisper decoder max target length
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--base_model", default="openai/whisper-small")
+    p.add_argument("--dataset", default="oddadmix/dialectal-arabic-lahgtna-v2-smaller-augmented")
+    p.add_argument("--language", default="ar")
+    p.add_argument("--run_name", default=None,
+                   help="Name for this run. Outputs go to runs/<run_name>/. "
+                        "Defaults to the base model name.")
+    p.add_argument("--output_dir", default=None,
+                   help="Override output dir (default: runs/<run_name>).")
+    p.add_argument("--notes", default="",
+                   help="Free-text note recorded in the run summary / README.")
+    p.add_argument("--per_device_train_batch_size", type=int, default=16)
+    p.add_argument("--per_device_eval_batch_size", type=int, default=8)
+    p.add_argument("--gradient_accumulation_steps", type=int, default=2)
+    p.add_argument("--learning_rate", type=float, default=1e-5)
+    p.add_argument("--warmup_steps", type=int, default=500)
+    p.add_argument("--max_steps", type=int, default=6000)
+    p.add_argument("--eval_steps", type=int, default=500)
+    p.add_argument("--save_steps", type=int, default=500)
+    p.add_argument("--num_workers", type=int, default=8)
+    p.add_argument("--normalize_letters", action="store_true",
+                   help="Also fold أإآ->ا, ى->ي, ة->ه (off by default).")
+    p.add_argument("--eval_only_frac", type=float, default=1.0,
+                   help="Use a fraction of the test split for periodic eval (speed).")
+    p.add_argument("--resume_from_checkpoint", default=None,
+                   help="Path to a checkpoint dir to resume training from.")
+    return p.parse_args()
+
+
+@dataclass
+class DataCollator:
+    """Extract log-mel features from raw audio and tokenize cleaned labels."""
+
+    processor: WhisperProcessor
+    normalize_letters: bool
+    decoder_start_token_id: int
+
+    def __call__(self, batch: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+        fe = self.processor.feature_extractor
+        tok = self.processor.tokenizer
+
+        arrays = [ex["audio"]["array"] for ex in batch]
+        feats = fe(arrays, sampling_rate=16000, return_tensors="pt")
+        out = {"input_features": feats.input_features}
+
+        texts = [clean_text(ex["text"], self.normalize_letters) for ex in batch]
+        # tok(texts) prepends the Whisper prefix (<|sot|><|ar|><|transcribe|>
+        # <|notimestamps|>) and appends <|eot|>. The Trainer re-prepends the
+        # decoder-start token, so we strip the leading <|sot|> below.
+        # truncation is a safety net; over-length rows are filtered out in main()
+        label_ids = tok(texts, max_length=MAX_LABEL_TOKENS, truncation=True).input_ids
+        labels = tok.pad({"input_ids": label_ids}, return_tensors="pt")
+        # mask padding so it's ignored by the loss
+        mask = labels.attention_mask.ne(1)
+        labels_ids = labels.input_ids.masked_fill(mask, -100)
+        # The tokenizer prepends <|startoftranscript|> (== decoder_start_token_id).
+        # The Trainer re-prepends it when building decoder_input_ids, so strip it
+        # here to avoid a doubled start token. NOTE: Whisper's bos_token_id is
+        # <|endoftext|> (50257), NOT the sot token, so we must compare against
+        # decoder_start_token_id explicitly.
+        if (labels_ids[:, 0] == self.decoder_start_token_id).all().cpu().item():
+            labels_ids = labels_ids[:, 1:]
+        out["labels"] = labels_ids
+        return out
+
+
+def build_metrics(processor, normalize_letters):
+    tok = processor.tokenizer
+
+    def compute_metrics(pred):
+        pred_ids = pred.predictions
+        label_ids = pred.label_ids
+        label_ids[label_ids == -100] = tok.pad_token_id
+
+        pred_str = tok.batch_decode(pred_ids, skip_special_tokens=True)
+        ref_str = tok.batch_decode(label_ids, skip_special_tokens=True)
+
+        preds = [clean_text(p, normalize_letters) for p in pred_str]
+        refs = [clean_text(r, normalize_letters) for r in ref_str]
+        # jiwer needs non-empty references; drop any degenerate pairs
+        pairs = [(p, r) for p, r in zip(preds, refs) if r.strip()]
+        if not pairs:
+            return {"wer": 1.0, "cer": 1.0}
+        preds, refs = map(list, zip(*pairs))
+        return {
+            "wer": jiwer.wer(refs, preds),
+            "cer": jiwer.cer(refs, preds),
+        }
+
+    return compute_metrics
+
+
+def keep_row(text: str, duration: float, normalize_letters: bool) -> bool:
+    if duration is None or not (MIN_AUDIO_SECONDS <= duration <= MAX_AUDIO_SECONDS):
+        return False
+    return bool(clean_text(text, normalize_letters))
+
+
+def main() -> None:
+    args = parse_args()
+
+    run_name = args.run_name or Path(args.base_model).name
+    output_dir = args.output_dir or f"runs/{run_name}"
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    print(f"run '{run_name}' -> {output_dir}")
+
+    processor = WhisperProcessor.from_pretrained(
+        args.base_model, language=args.language, task="transcribe"
+    )
+
+    ds = load_dataset(args.dataset)
+    # Filter on text+duration only -> no audio decode during filtering.
+    nl = args.normalize_letters
+    ds = ds.filter(
+        lambda text, duration: keep_row(text, duration, nl),
+        input_columns=["text", "duration"],
+        num_proc=args.num_workers,
+    )
+    # Drop rows whose (cleaned) transcript exceeds Whisper's 448-token target
+    # limit. These are corrupt/mislabeled clips (e.g. a huge text blob on a
+    # short clip) and would crash training. Tokenizing text is cheap; keep it
+    # single-process to avoid fast-tokenizer fork deadlocks.
+    tok = processor.tokenizer
+    before = {split: len(ds[split]) for split in ds}
+    ds = ds.filter(
+        lambda text: len(tok(clean_text(text, nl)).input_ids) <= MAX_LABEL_TOKENS,
+        input_columns=["text"],
+        num_proc=1,
+    )
+    after = {split: len(ds[split]) for split in ds}
+    print(f"kept {after} (dropped over-length: "
+          f"{ {s: before[s] - after[s] for s in before} })")
+
+    eval_ds = ds["test"]
+    if args.eval_only_frac < 1.0:
+        n = max(1, int(len(eval_ds) * args.eval_only_frac))
+        eval_ds = eval_ds.select(range(n))
+
+    # Force fp32 load: some checkpoints (e.g. large-v3-turbo) ship in fp16, which
+    # crashes generate() at eval (no autocast -> fp16 weights vs fp32 features).
+    # bf16-mixed training autocasts from fp32 weights, same as small/medium.
+    model = WhisperForConditionalGeneration.from_pretrained(
+        args.base_model, torch_dtype=torch.float32)
+    # Standard Whisper fine-tuning: let the model learn the language, don't
+    # force decoder ids or suppress tokens during training.
+    model.config.forced_decoder_ids = None
+    model.config.suppress_tokens = []
+    model.generation_config.language = args.language
+    model.generation_config.task = "transcribe"
+    model.generation_config.forced_decoder_ids = None
+    model.config.use_cache = False  # required with gradient checkpointing
+
+    collator = DataCollator(
+        processor=processor,
+        normalize_letters=nl,
+        decoder_start_token_id=model.config.decoder_start_token_id,
+    )
+
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=output_dir,
+        run_name=run_name,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        per_device_eval_batch_size=args.per_device_eval_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        warmup_steps=args.warmup_steps,
+        max_steps=args.max_steps,
+        gradient_checkpointing=True,
+        bf16=True,
+        fp16=False,
+        eval_strategy="steps",
+        eval_steps=args.eval_steps,
+        save_steps=args.save_steps,
+        logging_steps=25,
+        report_to=["tensorboard"],
+        predict_with_generate=True,
+        generation_max_length=225,
+        save_total_limit=3,
+        load_best_model_at_end=True,
+        metric_for_best_model="wer",
+        greater_is_better=False,
+        dataloader_num_workers=args.num_workers,
+        remove_unused_columns=False,   # collator needs the raw 'audio' column
+        label_names=["labels"],
+    )
+
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=ds["train"],
+        eval_dataset=eval_ds,
+        data_collator=collator,
+        compute_metrics=build_metrics(processor, nl),
+        processing_class=processor,
+    )
+
+    trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
+    trainer.save_model(output_dir)
+    processor.save_pretrained(output_dir)
+
+    # --- versioned run summary -------------------------------------------
+    evals = [h for h in trainer.state.log_history if "eval_wer" in h]
+    best = min(evals, key=lambda h: h["eval_wer"]) if evals else {}
+    train_logs = [h for h in trainer.state.log_history if "train_runtime" in h]
+    summary = {
+        "run_name": run_name,
+        "base_model": args.base_model,
+        "dataset": args.dataset,
+        "language": args.language,
+        "output_dir": output_dir,
+        "notes": args.notes,
+        "hyperparams": {
+            "learning_rate": args.learning_rate,
+            "warmup_steps": args.warmup_steps,
+            "max_steps": args.max_steps,
+            "per_device_train_batch_size": args.per_device_train_batch_size,
+            "gradient_accumulation_steps": args.gradient_accumulation_steps,
+            "effective_batch_size": args.per_device_train_batch_size
+            * args.gradient_accumulation_steps,
+            "normalize_letters": args.normalize_letters,
+        },
+        "train_examples": len(ds["train"]),
+        "eval_examples": len(eval_ds),
+        "best_wer": round(best.get("eval_wer", float("nan")), 4),
+        "best_cer": round(best.get("eval_cer", float("nan")), 4),
+        "best_step": best.get("step"),
+        "best_epoch": round(best.get("epoch", 0), 2),
+        "best_checkpoint": trainer.state.best_model_checkpoint,
+        "train_runtime_sec": round(train_logs[-1]["train_runtime"]) if train_logs else None,
+        "finished_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC"),
+        "eval_history": [
+            {"step": h["step"], "wer": round(h["eval_wer"], 4), "cer": round(h["eval_cer"], 4)}
+            for h in evals
+        ],
+    }
+    summary_path = Path(output_dir) / "summary.json"
+    summary_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False))
+    print(f"Done. Model saved to {output_dir}")
+    print(f"Best WER {summary['best_wer']} / CER {summary['best_cer']} "
+          f"@ step {summary['best_step']}. Summary -> {summary_path}")
+
+    # refresh the README run log (best-effort)
+    try:
+        import log_runs
+        log_runs.update_readme()
+        print("README run log updated.")
+    except Exception as e:  # never fail training over docs
+        print(f"(README auto-update skipped: {e})")
+
+
+if __name__ == "__main__":
+    main()