Create distill_smol-tts.py

distill_smol-tts.py

@@ -0,0 +1,272 @@
+#!/usr/bin/env python
+"""
+Colab-ready Knowledge Distillation script
+Teacher  : maya-research/veena-tts (4-bit)
+Student  : HuggingFaceTB/SmolLM-135M
+Dataset  : ArunKr/tts-hindi (HF Hub)
+Output   : push to hub -> ArunKr/smol-tts
+
+How to run on Colab (single GPU):
+-------------------------------------------------
+1) !apt-get -y install libsndfile1
+2) !pip install -U torch torchvision torchaudio transformers accelerate bitsandbytes snac soundfile wandb datasets einops sentencepiece
+3) !huggingface-cli login  # paste your token with write access to ArunKr
+4) Save this file as distill_colab.py and run:
+   !accelerate launch distill_colab.py \
+       --hf_dataset ArunKr/tts-hindi \
+       --split train \
+       --text_key text \
+       --output_dir /content/outputs \
+       --push_to_hub True \
+       --hub_model_id ArunKr/smol-tts
+"""
+import os
+import json
+import math
+import random
+import argparse
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from accelerate import Accelerator
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    get_cosine_schedule_with_warmup,
+)
+from datasets import load_dataset
+from huggingface_hub import HfApi
+
+# -----------------------------
+# Teacher code from the prompt
+# -----------------------------
+from snac import SNAC  # not used for KD but imported to satisfy trust_remote_code deps
+import soundfile as sf  # idem
+
+START_OF_SPEECH_TOKEN = 128257
+END_OF_SPEECH_TOKEN   = 128258
+START_OF_HUMAN_TOKEN  = 128259
+END_OF_HUMAN_TOKEN    = 128260
+START_OF_AI_TOKEN     = 128261
+END_OF_AI_TOKEN       = 128262
+AUDIO_CODE_BASE_OFFSET = 128266
+
+SPEAKERS = ["kavya", "agastya", "maitri", "vinaya"]
+
+
+def load_teacher(device_map="auto"):
+    quant_cfg = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+    )
+    teacher = AutoModelForCausalLM.from_pretrained(
+        "maya-research/veena-tts",
+        quantization_config=quant_cfg,
+        device_map=device_map,
+        trust_remote_code=True,
+    )
+    teacher_tok = AutoTokenizer.from_pretrained("maya-research/veena-tts", trust_remote_code=True)
+    return teacher, teacher_tok
+
+
+# -----------------------------
+# HF dataset wrapper
+# -----------------------------
+class HFDataset(Dataset):
+    def __init__(self, hf_ds, text_key, tokenizer, max_len):
+        self.ds = hf_ds
+        self.key = text_key
+        self.tok = tokenizer
+        self.max_len = max_len
+        if self.key not in self.ds.features:
+            raise ValueError(f"Column '{self.key}' not in dataset columns: {self.ds.features}")
+
+    def __len__(self):
+        return len(self.ds)
+
+    def __getitem__(self, idx):
+        text = self.ds[idx][self.key]
+        enc = self.tok(text, truncation=True, max_length=self.max_len, return_tensors="pt", add_special_tokens=True)
+        return {"input_ids": enc.input_ids[0], "attention_mask": enc.attention_mask[0], "text": text}
+
+
+def collate(batch, pad_id):
+    input_ids = [b["input_ids"] for b in batch]
+    attn = [b["attention_mask"] for b in batch]
+    maxlen = max(x.size(0) for x in input_ids)
+    def pad(x, val):
+        if x.size(0) == maxlen:
+            return x
+        return torch.cat([x, torch.full((maxlen - x.size(0),), val, dtype=x.dtype)], dim=0)
+    input_ids = torch.stack([pad(x, pad_id) for x in input_ids])
+    attn = torch.stack([pad(x, 0) for x in attn])
+    return {"input_ids": input_ids, "attention_mask": attn}
+
+
+# -----------------------------
+# Distillation utils
+# -----------------------------
+@torch.no_grad()
+def teacher_forward(teacher, teacher_tok, batch_text, device):
+    prompts = []
+    for txt in batch_text:
+        spk = random.choice(SPEAKERS)
+        p = f"<spk_{spk}> {txt}"
+        prompt_tokens = teacher_tok.encode(p, add_special_tokens=False)
+        seq = [START_OF_HUMAN_TOKEN, *prompt_tokens, END_OF_HUMAN_TOKEN, START_OF_AI_TOKEN]
+        prompts.append(seq)
+
+    max_len = max(len(p) for p in prompts)
+    input_ids = torch.full((len(prompts), max_len), teacher_tok.pad_token_id, dtype=torch.long, device=device)
+    attn_mask = torch.zeros_like(input_ids)
+    for i, seq in enumerate(prompts):
+        input_ids[i, : len(seq)] = torch.tensor(seq, device=device)
+        attn_mask[i, : len(seq)] = 1
+
+    out = teacher(input_ids=input_ids, attention_mask=attn_mask, output_logits=True)
+    return out.logits
+
+
+def kl_divergence(student_logits, teacher_logits, mask):
+    student_log_probs = F.log_softmax(student_logits, dim=-1)
+    teacher_probs = F.softmax(teacher_logits, dim=-1)
+    kl = F.kl_div(student_log_probs, teacher_probs, reduction='none').sum(-1)
+    kl = kl * mask
+    return kl.sum() / mask.sum().clamp(min=1)
+
+
+# -----------------------------
+# Main
+# -----------------------------
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--hf_dataset', type=str, default='ArunKr/tts-hindi')
+    parser.add_argument('--split', type=str, default='train')
+    parser.add_argument('--text_key', type=str, default='text')
+    parser.add_argument('--output_dir', type=str, required=True)
+    parser.add_argument('--epochs', type=int, default=3)
+    parser.add_argument('--batch_size', type=int, default=8)
+    parser.add_argument('--lr', type=float, default=2e-4)
+    parser.add_argument('--warmup_steps', type=int, default=500)
+    parser.add_argument('--max_len', type=int, default=512)
+    parser.add_argument('--seed', type=int, default=42)
+    parser.add_argument('--weight_ce', type=float, default=0.1)
+    parser.add_argument('--weight_kl', type=float, default=1.0)
+    parser.add_argument('--gradient_accumulation_steps', type=int, default=1)
+    parser.add_argument('--log_every', type=int, default=50)
+    parser.add_argument('--push_to_hub', type=lambda x: x.lower()=='true', default=False)
+    parser.add_argument('--hub_model_id', type=str, default='ArunKr/smol-tts')
+    args = parser.parse_args()
+
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    accelerator = Accelerator(log_with="wandb")
+    accelerator.init_trackers("distill-smollm135m", config=vars(args))
+
+    # Teacher
+    teacher, teacher_tok = load_teacher()
+    teacher.eval()
+
+    # Student
+    student_name = "HuggingFaceTB/SmolLM-135M"
+    student = AutoModelForCausalLM.from_pretrained(student_name)
+    student_tok = AutoTokenizer.from_pretrained(student_name)
+    if student_tok.pad_token is None:
+        student_tok.pad_token = student_tok.eos_token
+
+    # Dataset from HF Hub
+    raw_ds = load_dataset(args.hf_dataset, split=args.split)
+    ds = HFDataset(raw_ds, args.text_key, tokenizer=student_tok, max_len=args.max_len)
+    dl = DataLoader(ds, batch_size=args.batch_size, shuffle=True,
+                    collate_fn=lambda b: collate(b, pad_id=student_tok.pad_token_id))
+
+    # Optimizer & scheduler
+    optim = torch.optim.AdamW(student.parameters(), lr=args.lr)
+    total_steps = args.epochs * math.ceil(len(ds) / (args.batch_size * args.gradient_accumulation_steps))
+    sched = get_cosine_schedule_with_warmup(optim, num_warmup_steps=args.warmup_steps, num_training_steps=total_steps)
+
+    student, optim, dl, sched = accelerator.prepare(student, optim, dl, sched)
+
+    global_step = 0
+    for epoch in range(args.epochs):
+        for step, batch in enumerate(dl):
+            student.train()
+
+            # Get original texts back
+            pad_id = student_tok.pad_token_id
+            texts = []
+            for ids in batch['input_ids']:
+                valid = ids[ids != pad_id].tolist()
+                texts.append(student_tok.decode(valid, skip_special_tokens=True))
+
+            with torch.no_grad():
+                t_logits = teacher_forward(teacher, teacher_tok, texts, device=accelerator.device)
+
+            s_in = batch['input_ids']
+            attn = batch['attention_mask']
+            B, L = s_in.shape
+
+            # align teacher length
+            if t_logits.size(1) < L:
+                pad = torch.zeros(B, L - t_logits.size(1), t_logits.size(2), device=t_logits.device)
+                t_logits = torch.cat([t_logits, pad], dim=1)
+            elif t_logits.size(1) > L:
+                t_logits = t_logits[:, :L]
+
+            outputs = student(input_ids=s_in, attention_mask=attn)
+            s_logits = outputs.logits
+
+            loss_mask = attn.clone()
+            loss_mask[:, 0] = 0
+
+            kd_loss = kl_divergence(s_logits, t_logits, loss_mask) * args.weight_kl
+
+            shift_logits = s_logits[:, :-1].contiguous()
+            shift_labels = s_in[:, 1:].contiguous()
+            shift_mask = loss_mask[:, 1:]
+            ce_loss = torch.nn.functional.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                reduction='none'
+            )
+            ce_loss = (ce_loss.view(B, -1) * shift_mask).sum() / shift_mask.sum().clamp(min=1)
+            ce_loss = ce_loss * args.weight_ce
+
+            loss = (kd_loss + ce_loss) / args.gradient_accumulation_steps
+            accelerator.backward(loss)
+
+            if (step + 1) % args.gradient_accumulation_steps == 0:
+                optim.step(); sched.step(); optim.zero_grad(); global_step += 1
+
+            if accelerator.is_main_process and global_step % args.log_every == 0:
+                accelerator.log({"loss": loss.item(), "kd_loss": kd_loss.item(), "ce_loss": ce_loss.item(), "lr": sched.get_last_lr()[0]}, step=global_step)
+
+        if accelerator.is_main_process:
+            save_dir = Path(args.output_dir) / f"epoch_{epoch}"
+            save_dir.mkdir(parents=True, exist_ok=True)
+            accelerator.unwrap_model(student).save_pretrained(save_dir)
+            student_tok.save_pretrained(save_dir)
+
+    # Final save & (optional) push to hub
+    if accelerator.is_main_process:
+        final_dir = Path(args.output_dir) / "final"
+        final_dir.mkdir(parents=True, exist_ok=True)
+        accelerator.unwrap_model(student).save_pretrained(final_dir)
+        student_tok.save_pretrained(final_dir)
+        if args.push_to_hub:
+            from transformers import AutoModelForCausalLM
+            AutoModelForCausalLM.from_pretrained(final_dir).push_to_hub(args.hub_model_id)
+            student_tok.push_to_hub(args.hub_model_id)
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()