Upload run_finetune.py with huggingface_hub

run_finetune.py

@@ -0,0 +1,312 @@
+"""
+MonSub Whisper Full Fine-tune — RunPod A100 80GB.
+
+Base: Tsedee/whisper-large-v2-mn-monsub (existing Mongolian fine-tune)
+Data: 27 hours (4 HF datasets combined)
+Method: Full fine-tune (all 1.5B params)
+Output: Tsedee/whisper-large-v2-mn-monsub-v2
+
+Proven techniques:
+- Mixed precision bf16
+- Gradient checkpointing (saves VRAM)
+- Linear warmup + cosine decay
+- Mongolian forced_decoder_ids
+- Eval every 500 steps, save best by WER
+"""
+import os
+import sys
+import torch
+from dataclasses import dataclass
+from datasets import load_dataset, concatenate_datasets, Audio
+from transformers import (
+    WhisperForConditionalGeneration,
+    WhisperProcessor,
+    WhisperFeatureExtractor,
+    WhisperTokenizer,
+    Seq2SeqTrainingArguments,
+    Seq2SeqTrainer,
+    GenerationConfig,
+)
+from transformers.models.whisper.english_normalizer import BasicTextNormalizer
+import evaluate
+
+# ═══════════════════════════════════════════════════════════
+# CONFIG
+# ═══════════════════════════════════════════════════════════
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+BASE_MODEL = "Tsedee/whisper-large-v2-mn-monsub"
+OUTPUT_MODEL = "Tsedee/whisper-large-v2-mn-monsub-v2"
+OUTPUT_DIR = "/workspace/monsub-finetune"
+
+# Training hyperparams (proven for Whisper full fine-tune)
+BATCH_SIZE = 8           # A40 48GB
+GRAD_ACCUM = 4           # effective batch = 32
+LEARNING_RATE = 5e-6     # low LR for continued fine-tune (not from scratch)
+WARMUP_STEPS = 500
+NUM_EPOCHS = 3
+EVAL_STEPS = 500
+SAVE_STEPS = 500
+MAX_LABEL_LENGTH = 448
+LANGUAGE = "mn"
+TASK = "transcribe"
+
+# Datasets to combine
+DATASETS = [
+    {"name": "Tsedee/monsub-chimege-10h", "split": "train", "text_col": "sentence"},
+    {"name": "Tsedee/monsub-mongolian-asr", "split": "train", "text_col": "sentence"},
+    {"name": "Tsedee/mongolian-cv20-normalized", "split": "train", "text_col": "sentence"},
+    {"name": "Tsedee/monsub-chimege-youtube-9h", "split": "train", "text_col": "sentence"},
+]
+
+normalizer = BasicTextNormalizer()
+
+
+# ═══════════════════════════════════════════════════════════
+# DATA LOADING
+# ═══════════════════════════════════════════════════════════
+def load_all_datasets():
+    """Load and combine all datasets."""
+    print("=" * 60)
+    print("LOADING DATASETS")
+    print("=" * 60)
+
+    all_ds = []
+    total_hours = 0
+
+    for ds_info in DATASETS:
+        name = ds_info["name"]
+        text_col = ds_info["text_col"]
+        print(f"\n  Loading {name}...", flush=True)
+        try:
+            ds = load_dataset(name, split=ds_info["split"], token=HF_TOKEN)
+            # Normalize column names
+            if text_col != "sentence" and text_col in ds.column_names:
+                ds = ds.rename_column(text_col, "sentence")
+            # Ensure audio column
+            if "audio" in ds.column_names:
+                ds = ds.cast_column("audio", Audio(sampling_rate=16000))
+
+            # Calculate duration
+            if "duration" in ds.column_names:
+                hours = sum(ds["duration"]) / 3600
+            else:
+                hours = len(ds) * 10 / 3600  # estimate ~10s per sample
+
+            total_hours += hours
+            print(f"    ✅ {len(ds)} samples, ~{hours:.1f}h", flush=True)
+            all_ds.append(ds)
+        except Exception as e:
+            print(f"    ❌ Failed: {e}", flush=True)
+
+    # Combine
+    combined = concatenate_datasets(all_ds)
+    print(f"\n  TOTAL: {len(combined)} samples, ~{total_hours:.1f} hours")
+
+    # Train/test split (95/5)
+    split = combined.train_test_split(test_size=0.05, seed=42)
+    print(f"  Train: {len(split['train'])}, Test: {len(split['test'])}")
+    return split["train"], split["test"]
+
+
+# ═══════════════════════════════════════════════════════════
+# DATA PROCESSING
+# ═══════════════════════════════════════════════════════════
+def prepare_dataset(batch, processor):
+    """Process a batch: audio → features, text → labels."""
+    audio = batch["audio"]
+    inputs = processor.feature_extractor(
+        audio["array"], sampling_rate=audio["sampling_rate"]
+    )
+    batch["input_features"] = inputs.input_features[0]
+
+    # Tokenize text
+    text = batch["sentence"]
+    if not text or len(text.strip()) < 1:
+        text = " "
+    batch["labels"] = processor.tokenizer(text).input_ids
+    return batch
+
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: any
+    decoder_start_token_id: int
+
+    def __call__(self, features):
+        input_features = [{"input_features": f["input_features"]} for f in features]
+        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+
+        label_features = [{"input_ids": f["labels"]} for f in features]
+        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
+
+        labels = labels_batch["input_ids"].masked_fill(
+            labels_batch.attention_mask.ne(1), -100
+        )
+
+        # Remove decoder_start_token_id from labels
+        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+        return batch
+
+
+# ═══════════════════════════════════════════════════════════
+# METRICS
+# ═══════════════════════════════════════════════════════════
+wer_metric = evaluate.load("wer")
+
+
+def compute_metrics(pred, tokenizer):
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # Replace -100 with pad
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    # Normalize
+    pred_str = [normalizer(p) for p in pred_str]
+    label_str = [normalizer(l) for l in label_str]
+
+    wer = 100 * wer_metric.compute(predictions=pred_str, references=label_str)
+    return {"wer": wer}
+
+
+# ═══════════════════════════════════════════════════════════
+# MAIN
+# ═══════════════════════════════════════════════════════════
+def main():
+    print("=" * 60)
+    print("MonSub Whisper Full Fine-tune")
+    print(f"Base: {BASE_MODEL}")
+    print(f"Output: {OUTPUT_MODEL}")
+    print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
+    print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f}GB" if torch.cuda.is_available() else "")
+    print("=" * 60)
+
+    # Load model + processor
+    print("\n📦 Loading model...", flush=True)
+    processor = WhisperProcessor.from_pretrained(BASE_MODEL, token=HF_TOKEN)
+    model = WhisperForConditionalGeneration.from_pretrained(
+        BASE_MODEL, token=HF_TOKEN, torch_dtype=torch.bfloat16
+    )
+
+    # Fix generation config from base whisper-large-v2
+    print("  Fixing generation_config...", flush=True)
+    base_gc = GenerationConfig.from_pretrained("openai/whisper-large-v2")
+    model.generation_config = base_gc
+
+    # Set Mongolian
+    model.generation_config.forced_decoder_ids = processor.get_decoder_prompt_ids(
+        language=LANGUAGE, task=TASK
+    )
+    model.config.forced_decoder_ids = model.generation_config.forced_decoder_ids
+    model.config.suppress_tokens = []
+
+    # Enable gradient checkpointing (saves VRAM)
+    model.config.use_cache = False
+    model.gradient_checkpointing_enable()
+
+    print(f"  Model params: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M", flush=True)
+
+    # Load data
+    train_ds, eval_ds = load_all_datasets()
+
+    # Process datasets
+    print("\n🔄 Processing datasets...", flush=True)
+    train_ds = train_ds.map(
+        lambda x: prepare_dataset(x, processor),
+        remove_columns=train_ds.column_names,
+        num_proc=4,
+    )
+    eval_ds = eval_ds.map(
+        lambda x: prepare_dataset(x, processor),
+        remove_columns=eval_ds.column_names,
+        num_proc=4,
+    )
+
+    # Filter too-long labels
+    train_ds = train_ds.filter(lambda x: len(x["labels"]) < MAX_LABEL_LENGTH)
+    eval_ds = eval_ds.filter(lambda x: len(x["labels"]) < MAX_LABEL_LENGTH)
+    print(f"  After filter: train={len(train_ds)}, eval={len(eval_ds)}", flush=True)
+
+    # Data collator
+    data_collator = DataCollatorSpeechSeq2SeqWithPadding(
+        processor=processor,
+        decoder_start_token_id=model.config.decoder_start_token_id,
+    )
+
+    # Training args
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=OUTPUT_DIR,
+        per_device_train_batch_size=BATCH_SIZE,
+        per_device_eval_batch_size=8,
+        gradient_accumulation_steps=GRAD_ACCUM,
+        learning_rate=LEARNING_RATE,
+        warmup_steps=WARMUP_STEPS,
+        num_train_epochs=NUM_EPOCHS,
+        bf16=True,
+        evaluation_strategy="steps",
+        eval_steps=EVAL_STEPS,
+        save_strategy="steps",
+        save_steps=SAVE_STEPS,
+        save_total_limit=3,
+        load_best_model_at_end=True,
+        metric_for_best_model="wer",
+        greater_is_better=False,
+        predict_with_generate=True,
+        generation_max_length=225,
+        logging_steps=50,
+        report_to="none",
+        dataloader_num_workers=4,
+        push_to_hub=False,
+        lr_scheduler_type="cosine",
+        weight_decay=0.01,
+        gradient_checkpointing=True,
+        remove_unused_columns=False,
+    )
+
+    # Trainer
+    trainer = Seq2SeqTrainer(
+        args=training_args,
+        model=model,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        data_collator=data_collator,
+        compute_metrics=lambda pred: compute_metrics(pred, processor.tokenizer),
+        tokenizer=processor.feature_extractor,
+    )
+
+    # Train!
+    print("\n🚀 TRAINING STARTED!", flush=True)
+    print(f"  Epochs: {NUM_EPOCHS}")
+    print(f"  Batch: {BATCH_SIZE} × {GRAD_ACCUM} = {BATCH_SIZE * GRAD_ACCUM}")
+    print(f"  LR: {LEARNING_RATE}")
+    print(f"  Eval every: {EVAL_STEPS} steps")
+    print("=" * 60, flush=True)
+
+    trainer.train()
+
+    # Save best model
+    print("\n💾 Saving best model...", flush=True)
+    trainer.save_model(f"{OUTPUT_DIR}/best")
+    processor.save_pretrained(f"{OUTPUT_DIR}/best")
+
+    # Upload to HuggingFace
+    print(f"\n📤 Uploading to {OUTPUT_MODEL}...", flush=True)
+    model.push_to_hub(OUTPUT_MODEL, token=HF_TOKEN, private=True)
+    processor.push_to_hub(OUTPUT_MODEL, token=HF_TOKEN, private=True)
+
+    # Also upload generation_config
+    model.generation_config.save_pretrained(f"{OUTPUT_DIR}/best")
+
+    print(f"\n{'=' * 60}")
+    print(f"✅ DONE! Model: https://huggingface.co/{OUTPUT_MODEL}")
+    print(f"{'=' * 60}")
+
+
+if __name__ == "__main__":
+    main()