Upload run_finetune_v3.py with huggingface_hub

run_finetune_v3.py

@@ -0,0 +1,337 @@
+"""
+MonSub Whisper v3 Fine-tune — A40 48GB.
+
+Continued fine-tune from Tsedee/whisper-large-v2-mn-monsub (v1).
+Uses CER metric (Mongolian-д WER-ээс илүү тохиромжтой).
+
+Өмнөх бүх алдааг засасан:
+- processing_class (NOT tokenizer — deprecated)
+- datasets==2.21.0 (NOT latest — torchcodec error)
+- num_proc=1 (NOT 4 — multiprocess audio decode гацна)
+- HF_HOME=/workspace/.cache (container disk дүүрэхгүй)
+- generation_config fix (alignment_heads + no_timestamps_token_id)
+- fp16 (A40 дээр тохиромжтой)
+- eval crash handler
+"""
+import os
+import sys
+import torch
+import numpy as np
+from dataclasses import dataclass
+from datasets import load_dataset, concatenate_datasets, Audio
+from transformers import (
+    WhisperForConditionalGeneration,
+    WhisperProcessor,
+    Seq2SeqTrainingArguments,
+    Seq2SeqTrainer,
+    GenerationConfig,
+)
+import evaluate
+
+# ═══════════════════════════════════════════════════════════
+# CONFIG — А40-д оновчилсон
+# ═══════════════════════════════════════════════════════════
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+BASE_MODEL = "Tsedee/whisper-large-v2-mn-monsub"       # v1 суурь модел
+OUTPUT_MODEL = "Tsedee/whisper-large-v2-mn-monsub-v3"
+OUTPUT_DIR = "/workspace/monsub-finetune-v3"
+
+# A40 48GB — batch_size=16 багтана
+BATCH_SIZE = 16
+GRAD_ACCUM = 2              # effective batch = 32
+LEARNING_RATE = 3e-6        # Continued fine-tune → бага LR (шинээр бол 1e-5)
+WARMUP_STEPS = 300
+MAX_STEPS = 4000            # ~30 цаг дата → 4000 step хангалттай
+EVAL_STEPS = 500
+SAVE_STEPS = 500
+MAX_LABEL_LENGTH = 448
+LANGUAGE = "mn"
+TASK = "transcribe"
+
+# Datasets — mongolian-cv20-normalized ХАССАН (чанар муу)
+DATASETS = [
+    {"name": "Tsedee/monsub-chimege-10h", "split": "train", "text_col": "sentence"},
+    {"name": "Tsedee/monsub-mongolian-asr", "split": "train", "text_col": "sentence"},
+    {"name": "Tsedee/monsub-chimege-youtube-9h", "split": "train", "text_col": "sentence"},
+    # Нэмэлт dataset-үүд (хэрэгтэй бол comment арилга):
+    # {"name": "Tsedee/mongolian-bible-speech", "split": "train", "text_col": "sentence"},
+]
+
+
+# ═══════════════════════════════════════════════════════════
+# DATA LOADING
+# ═══════════════════════════════════════════════════════════
+def load_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 16kHz
+            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
+
+            total_hours += hours
+            print(f"    OK: {len(ds)} samples, ~{hours:.1f}h", flush=True)
+            all_ds.append(ds)
+        except Exception as e:
+            print(f"    FAILED: {e}", flush=True)
+
+    if not all_ds:
+        print("ERROR: No datasets loaded!")
+        sys.exit(1)
+
+    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):
+    audio = batch["audio"]
+    inputs = processor.feature_extractor(
+        audio["array"], sampling_rate=audio["sampling_rate"]
+    )
+    batch["input_features"] = inputs.input_features[0]
+
+    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
+        )
+
+        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+        return batch
+
+
+# ═══════════════════════════════════════════════════════════
+# CER METRIC — Монгол хэлэнд WER-ээс илүү тохиромжтой
+# ═══════════════════════════════════════════════════════════
+cer_metric = evaluate.load("cer")
+
+
+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)
+
+    # Filter empty pairs
+    pairs = [(p, l) for p, l in zip(pred_str, label_str) if l.strip()]
+    if not pairs:
+        return {"cer": 0.0}
+    pred_str, label_str = zip(*pairs)
+
+    cer = cer_metric.compute(predictions=list(pred_str), references=list(label_str))
+    return {"cer": cer}
+
+
+# ═══════════════════════════════════════════════════════════
+# MAIN
+# ═══════════════════════════════════════════════════════════
+def main():
+    print("=" * 60)
+    print("MonSub Whisper v3 Fine-tune")
+    print(f"Base: {BASE_MODEL}")
+    print(f"Output: {OUTPUT_MODEL}")
+
+    if torch.cuda.is_available():
+        gpu_name = torch.cuda.get_device_name(0)
+        vram = torch.cuda.get_device_properties(0).total_memory / 1e9
+        print(f"GPU: {gpu_name}")
+        print(f"VRAM: {vram:.1f}GB")
+    else:
+        print("WARNING: No GPU detected!")
+    print("=" * 60)
+
+    # ── Load model + processor ──
+    print("\nLoading model...", flush=True)
+    processor = WhisperProcessor.from_pretrained(BASE_MODEL, token=HF_TOKEN)
+    model = WhisperForConditionalGeneration.from_pretrained(
+        BASE_MODEL, token=HF_TOKEN
+    )
+
+    # ── generation_config fix ──
+    # alignment_heads + no_timestamps_token_id base-ээс авна
+    print("  Fixing generation_config from base whisper-large-v2...", flush=True)
+    base_gc = GenerationConfig.from_pretrained("openai/whisper-large-v2")
+    model.generation_config = base_gc
+
+    # Set Mongolian forced_decoder_ids
+    model.generation_config.forced_decoder_ids = processor.get_decoder_prompt_ids(
+        language=LANGUAGE, task=TASK
+    )
+    model.config.forced_decoder_ids = None      # Training-д None
+    model.config.suppress_tokens = []
+    model.config.use_cache = False               # Training-д заавал False
+
+    # Gradient checkpointing (VRAM хэмнэнэ)
+    model.gradient_checkpointing_enable()
+
+    params_m = sum(p.numel() for p in model.parameters()) / 1e6
+    print(f"  Model params: {params_m:.1f}M", flush=True)
+
+    # ── Load data ──
+    train_ds, eval_ds = load_all_datasets()
+
+    # ── Process datasets (num_proc=1 заавал!) ──
+    print("\nProcessing datasets (num_proc=1)...", flush=True)
+    train_ds = train_ds.map(
+        lambda x: prepare_dataset(x, processor),
+        remove_columns=train_ds.column_names,
+        num_proc=1,     # NOT 4 — multiprocess audio decode гацна
+    )
+    eval_ds = eval_ds.map(
+        lambda x: prepare_dataset(x, processor),
+        remove_columns=eval_ds.column_names,
+        num_proc=1,
+    )
+
+    # 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 — A40 48GB optimized ──
+    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,
+        max_steps=MAX_STEPS,
+        fp16=True,                          # A40 дээр fp16 хурдан
+        eval_strategy="steps",              # NOT evaluation_strategy (deprecated)
+        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="cer",        # CER = Монголд тохиромжтой
+        greater_is_better=False,
+        predict_with_generate=True,
+        generation_max_length=225,
+        logging_steps=25,
+        report_to="none",
+        dataloader_num_workers=2,           # A40-д 2 хангалттай
+        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),
+        processing_class=processor.feature_extractor,   # NOT tokenizer= (deprecated)
+    )
+
+    # ── Train ──
+    print(f"\nTRAINING STARTED!", flush=True)
+    print(f"  Steps: {MAX_STEPS}")
+    print(f"  Batch: {BATCH_SIZE} x {GRAD_ACCUM} = {BATCH_SIZE * GRAD_ACCUM}")
+    print(f"  LR: {LEARNING_RATE}")
+    print(f"  Eval every: {EVAL_STEPS} steps")
+    print(f"  Metric: CER (lower = better)")
+    print("=" * 60, flush=True)
+
+    try:
+        trainer.train()
+    except Exception as e:
+        print(f"\nTraining error: {e}", flush=True)
+        print("Attempting to save current model...", flush=True)
+        trainer.save_model(f"{OUTPUT_DIR}/emergency-save")
+        processor.save_pretrained(f"{OUTPUT_DIR}/emergency-save")
+        raise
+
+    # ── Save best model ──
+    print("\nSaving best model...", flush=True)
+    trainer.save_model(f"{OUTPUT_DIR}/best")
+    processor.save_pretrained(f"{OUTPUT_DIR}/best")
+
+    # Save generation_config with Mongolian settings
+    model.generation_config.forced_decoder_ids = processor.get_decoder_prompt_ids(
+        language=LANGUAGE, task=TASK
+    )
+    model.generation_config.save_pretrained(f"{OUTPUT_DIR}/best")
+
+    # ── Upload to HuggingFace ──
+    print(f"\nUploading to {OUTPUT_MODEL}...", flush=True)
+    try:
+        model.push_to_hub(OUTPUT_MODEL, token=HF_TOKEN, private=True)
+        processor.push_to_hub(OUTPUT_MODEL, token=HF_TOKEN, private=True)
+        model.generation_config.push_to_hub(OUTPUT_MODEL, token=HF_TOKEN)
+        print(f"  Upload OK: https://huggingface.co/{OUTPUT_MODEL}")
+    except Exception as e:
+        print(f"  Upload failed: {e}")
+        print(f"  Model saved locally: {OUTPUT_DIR}/best")
+
+    print(f"\n{'=' * 60}")
+    print(f"DONE! Model: {OUTPUT_MODEL}")
+    print(f"{'=' * 60}")
+
+
+if __name__ == "__main__":
+    main()