Upload folder using huggingface_hub

train_v3.py

@@ -0,0 +1,288 @@
+"""
+Voxtral 3B — Haitian Creole LoRA v3.
+Synthesizes domain-specific audio on-the-fly, combines with CMU dataset, trains.
+All in one job — no separate upload step.
+"""
+
+import torch
+import json
+import os
+import time
+import numpy as np
+import soundfile as sf
+from datasets import load_dataset, Dataset, Audio, concatenate_datasets
+from transformers import (
+    VoxtralForConditionalGeneration,
+    VoxtralProcessor,
+    VitsModel,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+)
+from peft import LoraConfig, get_peft_model, PeftModel
+from jiwer import wer, cer
+
+MODEL_ID = "mistralai/Voxtral-Mini-3B-2507"
+OUTPUT = os.environ.get("HF_OUTPUT_DIR", "./output")
+NUM_TEST = 200
+MAX_STEPS = 5000
+LR = 3e-5
+LORA_R = 32
+LORA_ALPHA = 64
+SAMPLE_RATE = 16000
+
+
+class VoxtralDataCollator:
+    def __init__(self, processor, model_id):
+        self.processor = processor
+        self.model_id = model_id
+
+    def __call__(self, features):
+        texts = [f["text"].lower() for f in features]
+        audios = [f["audio"]["array"] for f in features]
+
+        prompt = self.processor.apply_transcription_request(
+            language="en", model_id=self.model_id,
+            audio=audios, format=["WAV"] * len(audios),
+            return_tensors="pt",
+        )
+        passthrough = {k: v for k, v in prompt.items()
+                       if k not in ("input_ids", "attention_mask")}
+        prompt_ids = prompt["input_ids"]
+        prompt_attn = prompt["attention_mask"]
+        B = prompt_ids.size(0)
+        tok = self.processor.tokenizer
+
+        text_tok = tok(texts, add_special_tokens=False, padding=False,
+                       truncation=True, max_length=256, return_tensors=None)
+        text_ids_list = text_tok["input_ids"]
+
+        input_ids, attention_mask, labels = [], [], []
+        for i in range(B):
+            p_ids = prompt_ids[i].tolist()
+            p_att = prompt_attn[i].tolist()
+            t_ids = text_ids_list[i]
+            ids = p_ids + t_ids + [tok.eos_token_id]
+            attn = p_att + [1] * (len(t_ids) + 1)
+            lab = [-100] * len(p_ids) + t_ids + [tok.eos_token_id]
+            input_ids.append(ids)
+            attention_mask.append(attn)
+            labels.append(lab)
+
+        pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
+        max_len = max(len(x) for x in input_ids)
+        def pad_to(seq, fill, L): return seq + [fill] * (L - len(seq))
+
+        input_ids = [pad_to(x, pad_id, max_len) for x in input_ids]
+        attention_mask = [pad_to(x, 0, max_len) for x in attention_mask]
+        labels = [pad_to(x, -100, max_len) for x in labels]
+
+        batch = {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+        }
+        for k, v in passthrough.items():
+            batch[k] = v
+        return batch
+
+
+def synthesize_domain_data():
+    """Synthesize Creole audio for domain text datasets using MMS-TTS."""
+    print("\n  Loading MMS-TTS (facebook/mms-tts-hat)...")
+    tts_model = VitsModel.from_pretrained("facebook/mms-tts-hat").to("cuda")
+    tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-hat")
+
+    domains = [
+        ("jsbeaudry/med-eng-cre", "medical", "text_target", "train"),
+        ("jsbeaudry/data-therapy-creole", "therapy", "output", "train"),
+        ("jsbeaudry/creole-psychology", "psychology", "response_j", "train"),
+        ("jsbeaudry/creole-economy-with-words", "economics", "assistant", "train"),
+        ("jsbeaudry/creole-sport", "sports", "answer", "train"),
+    ]
+
+    audio_dir = "/tmp/synth_audio"
+    os.makedirs(audio_dir, exist_ok=True)
+
+    all_audio_paths = []
+    all_texts = []
+
+    for ds_name, domain, text_col, split in domains:
+        print(f"  Synthesizing {domain}...")
+        ds = load_dataset(ds_name, split=split)
+        count = 0
+
+        for i, row in enumerate(ds):
+            text = row.get(text_col, "")
+            if not text or not isinstance(text, str) or len(text.strip()) < 10:
+                continue
+            text = text.strip()[:300]
+            try:
+                inputs = tts_tokenizer(text, return_tensors="pt").to("cuda")
+                with torch.no_grad():
+                    audio = tts_model(**inputs).waveform.squeeze().cpu().numpy()
+                if len(audio) > SAMPLE_RATE:  # >1 second
+                    path = os.path.join(audio_dir, f"{domain}_{i:05d}.wav")
+                    sf.write(path, audio, SAMPLE_RATE)
+                    all_audio_paths.append(path)
+                    all_texts.append(text.lower())
+                    count += 1
+            except Exception:
+                continue
+
+        print(f"    {domain}: {count} samples")
+
+    # Free TTS model
+    del tts_model, tts_tokenizer
+    torch.cuda.empty_cache()
+
+    # Build HF dataset
+    synth_ds = Dataset.from_dict({"audio": all_audio_paths, "text": all_texts})
+    synth_ds = synth_ds.cast_column("audio", Audio(sampling_rate=SAMPLE_RATE))
+    print(f"  Total synthesized: {len(synth_ds)}")
+    return synth_ds
+
+
+def main():
+    gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
+    vram = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0
+
+    print("=" * 70)
+    print("  Voxtral 3B — Haitian Creole LoRA v3 (CMU + Domain Synthesis)")
+    print(f"  GPU: {gpu} ({vram:.0f} GB)")
+    print("=" * 70)
+
+    # Step 1: Synthesize domain audio
+    print("\n[1/6] Synthesizing domain-specific audio...")
+    t_synth = time.time()
+    synth_ds = synthesize_domain_data()
+    synth_time = time.time() - t_synth
+    print(f"  Synthesis done: {synth_time/60:.1f} min")
+
+    # Step 2: Load CMU dataset
+    print("\n[2/6] Loading CMU Haitian Creole dataset...")
+    cmu_ds = load_dataset("jsbeaudry/cmu_haitian_creole_speech", split="train")
+    cmu_ds = cmu_ds.cast_column("audio", Audio(sampling_rate=SAMPLE_RATE))
+    cmu_ds = cmu_ds.remove_columns(["text"]).rename_column("normalized_text", "text")
+    cmu_ds = cmu_ds.select_columns(["audio", "text"])
+    print(f"  CMU: {len(cmu_ds)} samples")
+
+    # Step 3: Merge and split
+    merged = concatenate_datasets([cmu_ds, synth_ds])
+    merged = merged.shuffle(seed=42)
+    test_ds = merged.select(range(len(merged) - NUM_TEST, len(merged)))
+    train_ds = merged.select(range(len(merged) - NUM_TEST))
+    print(f"  Merged: {len(merged)} total, Train: {len(train_ds)}, Test: {len(test_ds)}")
+
+    # Step 4: Load model + LoRA
+    print("\n[3/6] Loading Voxtral 3B + LoRA...")
+    processor = VoxtralProcessor.from_pretrained(MODEL_ID)
+    model = VoxtralForConditionalGeneration.from_pretrained(
+        MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto"
+    )
+    config = LoraConfig(
+        r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=0.05, bias="none",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        task_type="SEQ_2_SEQ_LM",
+    )
+    model = get_peft_model(model, config)
+    model.print_trainable_parameters()
+
+    collator = VoxtralDataCollator(processor, MODEL_ID)
+
+    # Step 5: Train
+    print(f"\n[4/6] Training {MAX_STEPS} steps...")
+    t_train = time.time()
+
+    args = TrainingArguments(
+        output_dir=OUTPUT,
+        max_steps=MAX_STEPS,
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=4,
+        learning_rate=LR,
+        weight_decay=0.01,
+        warmup_steps=100,
+        logging_steps=100,
+        save_steps=1000,
+        bf16=True,
+        max_grad_norm=1.0,
+        dataloader_num_workers=1,
+        remove_unused_columns=False,
+        report_to="none",
+        lr_scheduler_type="cosine",
+    )
+
+    trainer = Trainer(
+        model=model, args=args,
+        train_dataset=train_ds,
+        data_collator=collator,
+    )
+    trainer.train()
+    train_time = time.time() - t_train
+    print(f"  Training done: {train_time/60:.1f} min")
+
+    trainer.save_model(OUTPUT)
+    processor.save_pretrained(OUTPUT)
+
+    # Step 6: Evaluate
+    print(f"\n[5/6] Evaluating on {NUM_TEST} held-out samples...")
+    model.config = model.base_model.model.config
+    refs, preds = [], []
+
+    for i, sample in enumerate(test_ds):
+        audio = sample["audio"]["array"]
+        ref = sample["text"].lower()
+        inputs = processor.apply_transcription_request(
+            language="en", model_id=MODEL_ID,
+            audio=[audio], format=["WAV"], return_tensors="pt",
+        )
+        inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+        with torch.no_grad():
+            out = model.generate(**inputs, max_new_tokens=256)
+        pred = processor.batch_decode(out, skip_special_tokens=True)[0].strip()
+        refs.append(ref)
+        preds.append(pred)
+        if i < 10 or (i + 1) % 50 == 0:
+            print(f"  [{i+1}/{NUM_TEST}]")
+            if i < 5:
+                print(f"    Ref:  {ref[:70]}")
+                print(f"    Pred: {pred[:70]}")
+
+    ft_wer = wer(refs, preds)
+    ft_cer = cer(refs, preds)
+    exact = sum(1 for r, p in zip(refs, preds) if r.strip() == p.strip())
+
+    print(f"\n{'=' * 70}")
+    print("  RESULTS v3 (CMU + Medical + Therapy + Psychology + Economics + Sports)")
+    print(f"{'=' * 70}")
+    print(f"  v1 WER:  11.8%  (CMU only, r=16, 2000 steps)")
+    print(f"  v2 WER:   9.7%  (CMU only, r=32, 5000 steps)")
+    print(f"  v3 WER:  {ft_wer:.1%}  (CMU + domains, r=32, {MAX_STEPS} steps)")
+    print(f"  v3 CER:  {ft_cer:.1%}")
+    print(f"  v3 Exact: {exact}/{NUM_TEST} ({exact/NUM_TEST:.1%})")
+    print(f"  Training data: {len(train_ds)} samples ({len(cmu_ds)} CMU + {len(synth_ds)} synthetic)")
+    print(f"  Synthesis: {synth_time/60:.1f} min, Training: {train_time/60:.1f} min")
+    print(f"{'=' * 70}")
+
+    results = {
+        "version": "v3",
+        "finetuned_wer": ft_wer, "finetuned_cer": ft_cer,
+        "exact_matches": exact, "exact_match_rate": exact / NUM_TEST,
+        "v1_wer": 0.118, "v2_wer": 0.097, "baseline_wer": 1.014,
+        "train_samples": len(train_ds), "cmu_samples": len(cmu_ds),
+        "synth_samples": len(synth_ds), "test_samples": NUM_TEST,
+        "training_min": round(train_time / 60, 1),
+        "synthesis_min": round(synth_time / 60, 1),
+        "steps": MAX_STEPS, "lora_r": LORA_R, "lr": LR,
+        "gpu": gpu,
+        "domains": ["medical", "therapy", "psychology", "economics", "sports"],
+        "samples": [{"ref": r, "pred": p} for r, p in zip(refs, preds)],
+    }
+    os.makedirs(OUTPUT, exist_ok=True)
+    with open(os.path.join(OUTPUT, "results_v3.json"), "w") as f:
+        json.dump(results, f, indent=2, ensure_ascii=False)
+    print(f"Saved to {OUTPUT}/results_v3.json")
+
+
+if __name__ == "__main__":
+    main()