Add trainer script itself

whisper_medium.py

@@ -0,0 +1,249 @@
+# ============================================================
+# RESUMABLE WHISPER TRAINING SCRIPT WITH TIMESTAMP SUPPORT
+# ============================================================
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+import os
+import gc
+import torch
+import pandas as pd
+import evaluate
+
+from datasets import Dataset, Audio
+from transformers import (
+    WhisperForConditionalGeneration,
+    WhisperProcessor,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+)
+
+# ============================================================
+# CONFIG
+# ============================================================
+
+MODEL_SIZE = "medium"
+BASE_MODEL = f"neurlang/ipa-whisper-{MODEL_SIZE}"
+OUTPUT_DIR = f"whisper-{MODEL_SIZE}-finetuned"
+
+RESUME_TRAINING = False        # 🔁 flip to True to resume
+RESUME_CHECKPOINT = "checkpoint-1840000"       # e.g. "checkpoint-40000"
+RESUME_CHECKPOINT_TARGET = 1880000             # e.g.             80000
+
+# don't tune this, it's auto tuned on training start/resume
+lr = 1.251564455569462e-07 # 1e-5
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ============================================================
+# LOAD DATA
+#
+# FORMAT:
+# foo.mp3,hello world
+#
+# ============================================================
+
+train_df = pd.read_csv("train.csv")
+eval_df = pd.read_csv("test.csv")
+
+train_df.columns = ["audio", "sentence"]
+eval_df.columns = ["audio", "sentence"]
+
+train_dataset = Dataset.from_pandas(train_df)
+eval_dataset = Dataset.from_pandas(eval_df)
+
+train_dataset = train_dataset.cast_column("audio", Audio(sampling_rate=16000))
+eval_dataset = eval_dataset.cast_column("audio", Audio(sampling_rate=16000))
+
+# Shuffle the dataset with exact seed control
+train_dataset = train_dataset.shuffle(seed=42)  # Default shuffles all
+
+# ============================================================
+# PROCESSOR (TOKENIZER + FEATURE EXTRACTOR)
+# ============================================================
+
+if RESUME_TRAINING:
+    print(f"🔁 Loading processor from {OUTPUT_DIR}")
+    processor = WhisperProcessor.from_pretrained(OUTPUT_DIR)
+else:
+    print("🆕 Creating new processor")
+    processor = WhisperProcessor.from_pretrained(
+        BASE_MODEL,
+        language="english",
+        task="transcribe",
+        predict_timestamps=True,
+    )
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    processor.save_pretrained(OUTPUT_DIR)  # 🔒 critical
+
+# ============================================================
+# DATA PREPARATION
+# ============================================================
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+
+    batch["input_features"] = processor.feature_extractor(
+        audio["array"],
+        sampling_rate=16000
+    ).input_features[0]
+
+    text = batch["sentence"] if batch["sentence"] else ""
+    batch["labels"] = processor.tokenizer(
+        text,
+        return_tensors="pt"
+    ).input_ids[0]
+
+    del batch["audio"]
+    del batch["sentence"]
+    return batch
+
+train_dataset = train_dataset.map(prepare_dataset, num_proc=1)
+eval_dataset = eval_dataset.map(prepare_dataset, num_proc=1)
+
+# ============================================================
+# DATA COLLATOR
+# ============================================================
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: Any
+
+    def __call__(self, features):
+        inputs = [{"input_features": f["input_features"]} for f in features]
+        batch = self.processor.feature_extractor.pad(
+            inputs, return_tensors="pt"
+        )
+
+        labels = [{"input_ids": f["labels"]} for f in features]
+        labels_batch = self.processor.tokenizer.pad(
+            labels, return_tensors="pt"
+        )
+
+        labels = labels_batch["input_ids"].masked_fill(
+            labels_batch.attention_mask.ne(1), -100
+        )
+
+        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+        return batch
+
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor)
+
+# ============================================================
+# METRICS
+# ============================================================
+
+cer_metric = evaluate.load("cer")
+
+def compute_metrics(pred):
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    label_ids[label_ids == -100] = processor.tokenizer.pad_token_id
+
+    pred_str = processor.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = processor.tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    return {
+        "cer": 100 * cer_metric.compute(
+            predictions=pred_str,
+            references=label_str
+        )
+    }
+
+# ============================================================
+# TRAINING ARGUMENTS
+# ============================================================
+
+
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir=OUTPUT_DIR,
+    per_device_train_batch_size=4,
+    per_device_eval_batch_size=1,
+    learning_rate=lr,
+    warmup_steps=1000,
+    max_steps=RESUME_CHECKPOINT_TARGET,
+    evaluation_strategy="steps",
+    save_strategy="steps",
+    logging_steps=10*100,
+    eval_steps=10*100,
+    save_steps=10*100,
+    save_total_limit=3,
+    predict_with_generate=True,
+    generation_max_length=225,
+    fp16=False,
+    report_to=["tensorboard"],
+    load_best_model_at_end=False,
+    metric_for_best_model="cer",
+    greater_is_better=False,
+    save_safetensors=True,   # 🔒 ensure safetensors
+)
+
+# ============================================================
+# LOAD MODEL
+# ============================================================
+
+if RESUME_TRAINING:
+    assert RESUME_CHECKPOINT is not None, "RESUME_CHECKPOINT must be set"
+
+    checkpoint_path = os.path.join(OUTPUT_DIR, RESUME_CHECKPOINT)
+    print(f"🔁 Loading model from {checkpoint_path}")
+
+    model = WhisperForConditionalGeneration.from_pretrained(
+        checkpoint_path,
+        torch_dtype=torch.float32,
+    )
+else:
+    print("🆕 Loading base model")
+    model = WhisperForConditionalGeneration.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=torch.float32,
+    )
+
+# 🔒 Modified safety check for Transformers version
+# Transformers' Whisper uses different parameter naming
+print(f"✅ Model loaded successfully")
+print(f"   Model type: {type(model)}")
+print(f"   Device: {next(model.parameters()).device}")
+
+model.to(DEVICE)
+
+# ============================================================
+# TRAINER
+# ============================================================
+
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+)
+
+# ============================================================
+# TRAIN
+# ============================================================
+
+torch.cuda.empty_cache()
+gc.collect()
+
+if RESUME_TRAINING:
+    trainer.train(resume_from_checkpoint=checkpoint_path)
+else:
+    trainer.train()
+
+# ============================================================
+# SAVE FINAL
+# ============================================================
+
+trainer.save_model(OUTPUT_DIR)
+processor.save_pretrained(OUTPUT_DIR)
+
+print("✅ Training completed successfully")
+