ai-agent-book/chapter8/speech-sft-experiment/run_orpheus.py

#!/usr/bin/env python3
"""Run the Orpheus half of Experiment 8-6 on one local CUDA GPU.

The campaign deliberately keeps a held-out split and emits base/adapted audio
for identical prompts and seeds.  It is bounded for a workstation, but it is
not a one-batch smoke test: the default run encodes 144 real utterances and
performs 60 optimizer updates with effective batch size four.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import time
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
import torchaudio.functional as AF
from datasets import load_dataset
from huggingface_hub import HfApi
from snac import SNAC
from unsloth import FastLanguageModel
from transformers import Trainer, TrainingArguments

BASE_MODEL = "unsloth/orpheus-3b-0.1-ft"
DATASET = "maxbsoft/mrdragonfox-elise"
DATASET_REVISION = "2cc657c3f94a83df18fcd968b7531ca1a19c7f88"
SEED = 7601

PROMPTS = [
    "The morning train crossed the bridge just before sunrise.",
    "Please leave the blue notebook beside the kitchen window.",
    "A patient astronomer mapped every bright star in the winter sky.",
    "We walked home slowly while the last shops turned off their lights.",
    "Could you read the final paragraph one more time for the group?",
    "The small garden stayed green even through the hottest week of July.",
    "I packed a warm coat, two apples, and a compass for the long hike.",
    "Tomorrow's meeting begins at nine, so I will arrive a little early.",
]


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


class PadCollator:
    def __init__(self, pad_id: int):
        self.pad_id = pad_id

    def __call__(self, rows):
        n = max(len(x["input_ids"]) for x in rows)
        ids, labels, masks = [], [], []
        for row in rows:
            d = n - len(row["input_ids"])
            ids.append(row["input_ids"] + [self.pad_id] * d)
            labels.append(row["labels"] + [-100] * d)
            masks.append(row["attention_mask"] + [0] * d)
        return {
            "input_ids": torch.tensor(ids, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
            "attention_mask": torch.tensor(masks, dtype=torch.long),
        }


def encode_audio(snac, wave, sample_rate: int, seconds: float) -> list[int]:
    wave = torch.as_tensor(np.asarray(wave), dtype=torch.float32)
    if sample_rate != 24000:
        wave = AF.resample(wave, sample_rate, 24000)
    wave = wave[: int(seconds * 24000)]
    if wave.numel() < 2400:
        raise ValueError("utterance is shorter than 100 ms")
    with torch.inference_mode():
        codes = snac.encode(wave[None, None].cuda())
    out = []
    for i in range(codes[0].shape[1]):
        out.extend(
            [
                codes[0][0][i].item() + 128266,
                codes[1][0][2 * i].item() + 128266 + 4096,
                codes[2][0][4 * i].item() + 128266 + 2 * 4096,
                codes[2][0][4 * i + 1].item() + 128266 + 3 * 4096,
                codes[1][0][2 * i + 1].item() + 128266 + 4 * 4096,
                codes[2][0][4 * i + 2].item() + 128266 + 5 * 4096,
                codes[2][0][4 * i + 3].item() + 128266 + 6 * 4096,
            ]
        )
    # Remove codec frames whose first code repeats, matching the upstream recipe.
    dedup = out[:7]
    for i in range(7, len(out), 7):
        if out[i] != dedup[-7]:
            dedup.extend(out[i : i + 7])
    return dedup


def prepare_rows(ds, tokenizer, snac, indices, seconds):
    result, failures = [], []
    for pos, idx in enumerate(indices, 1):
        try:
            row = ds[int(idx)]
            codes = encode_audio(
                snac, row["audio"]["array"], row["audio"]["sampling_rate"], seconds
            )
            text_ids = tokenizer.encode(row["text"], add_special_tokens=True) + [128009]
            ids = [128259] + text_ids + [128260, 128261, 128257] + codes + [128258, 128262]
            result.append({"input_ids": ids, "labels": ids.copy(), "attention_mask": [1] * len(ids)})
        except Exception as exc:  # retained in the manifest
            failures.append({"dataset_index": int(idx), "error": repr(exc)})
        print(f"encoded {pos}/{len(indices)}", flush=True)
    return result, failures


def decode(snac, ids):
    speech = (ids == 128257).nonzero(as_tuple=True)[0]
    row = ids[speech[-1].item() + 1 :] if speech.numel() else ids
    eos = (row == 128258).nonzero(as_tuple=True)[0]
    if eos.numel():
        row = row[: eos[0].item()]
    values = [int(x) - 128266 for x in row[: (len(row) // 7) * 7]]
    layers = [[], [], []]
    invalid_frame = None
    for i in range(len(values) // 7):
        c = [values[7 * i + j] - j * 4096 for j in range(7)]
        if any(x < 0 or x > 4095 for x in c):
            invalid_frame = i
            break
        layers[0].append(c[0])
        layers[1].extend([c[1], c[4]])
        layers[2].extend([c[2], c[3], c[5], c[6]])
    if not layers[0]:
        return np.zeros(2400, dtype=np.float32), invalid_frame, 0
    tensors = [torch.tensor(x, dtype=torch.long)[None] for x in layers]
    with torch.inference_mode():
        audio = snac.cpu().decode(tensors).squeeze().float().numpy()
    return audio, invalid_frame, len(layers[0])


def generate_arm(model, tokenizer, snac, arm: str, out: Path, max_tokens: int):
    arm_dir = out / "audio" / "orpheus" / arm
    arm_dir.mkdir(parents=True, exist_ok=True)
    rows = []
    FastLanguageModel.for_inference(model)
    for i, prompt in enumerate(PROMPTS):
        torch.manual_seed(SEED + i)
        ids = tokenizer(prompt, return_tensors="pt").input_ids
        ids = torch.cat([torch.tensor([[128259]]), ids, torch.tensor([[128009, 128260]])], dim=1).cuda()
        with torch.inference_mode():
            generated = model.generate(
                input_ids=ids,
                attention_mask=torch.ones_like(ids),
                max_new_tokens=max_tokens,
                do_sample=True,
                temperature=0.6,
                top_p=0.95,
                repetition_penalty=1.1,
                eos_token_id=128258,
                use_cache=True,
            )[0].cpu()
        audio, invalid_frame, frames = decode(snac, generated)
        path = arm_dir / f"prompt_{i:02d}.wav"
        sf.write(path, audio, 24000, subtype="PCM_16")
        rows.append(
            {
                "prompt_id": i,
                "prompt": prompt,
                "seed": SEED + i,
                "path": str(path.relative_to(out)),
                "sha256": sha256(path),
                "samples": int(len(audio)),
                "seconds": len(audio) / 24000,
                "decoded_frames": frames,
                "first_invalid_frame": invalid_frame,
            }
        )
        print(f"generated Orpheus {arm} {i + 1}/{len(PROMPTS)}", flush=True)
    return rows


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--train-examples", type=int, default=128)
    p.add_argument("--eval-examples", type=int, default=16)
    p.add_argument("--max-audio-seconds", type=float, default=4.0)
    p.add_argument("--steps", type=int, default=60)
    p.add_argument("--generation-tokens", type=int, default=560)
    p.add_argument("--hf-repo", default="bojieli/exp8-6-orpheus-elise-lora")
    args = p.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    started = time.time()
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)

    ds = load_dataset(DATASET, revision=DATASET_REVISION, split="train")
    candidates = [i for i, x in enumerate(ds["duration"]) if 1.0 <= float(x) <= 10.5]
    random.shuffle(candidates)
    selected = candidates[: args.train_examples + args.eval_examples]

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=BASE_MODEL, max_seq_length=3072, dtype=None, load_in_4bit=False
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_alpha=16,
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=SEED,
    )
    snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").cuda().eval()
    train_rows, train_failures = prepare_rows(
        ds, tokenizer, snac, selected[: args.train_examples], args.max_audio_seconds
    )
    eval_rows, eval_failures = prepare_rows(
        ds, tokenizer, snac, selected[args.train_examples :], args.max_audio_seconds
    )
    snac.cpu()
    torch.cuda.empty_cache()

    base_audio = generate_arm(model, tokenizer, snac, "base", args.output, args.generation_tokens)
    FastLanguageModel.for_training(model)
    trainer = Trainer(
        model=model,
        train_dataset=train_rows,
        eval_dataset=eval_rows,
        data_collator=PadCollator(128263),
        args=TrainingArguments(
            output_dir=str(args.output / "orpheus_checkpoints"),
            per_device_train_batch_size=1,
            per_device_eval_batch_size=1,
            gradient_accumulation_steps=4,
            max_steps=args.steps,
            warmup_steps=5,
            learning_rate=2e-4,
            bf16=True,
            logging_steps=1,
            eval_strategy="no",
            save_strategy="no",
            optim="adamw_8bit",
            weight_decay=0.01,
            lr_scheduler_type="linear",
            seed=SEED,
            report_to="none",
        ),
    )
    pre_eval = trainer.evaluate()
    train_result = trainer.train()
    post_eval = trainer.evaluate()
    adapter_dir = args.output / "adapters" / "orpheus"
    model.save_pretrained(adapter_dir)
    tokenizer.save_pretrained(adapter_dir)
    model.push_to_hub(args.hf_repo, private=False, token=os.environ.get("HF_TOKEN"))
    tokenizer.push_to_hub(args.hf_repo, private=False, token=os.environ.get("HF_TOKEN"))
    adapter_revision = HfApi().model_info(args.hf_repo).sha
    adapted_audio = generate_arm(model, tokenizer, snac, "adapted", args.output, args.generation_tokens)

    adapter_files = [
        {"path": str(x.relative_to(args.output)), "bytes": x.stat().st_size, "sha256": sha256(x)}
        for x in sorted(adapter_dir.rglob("*"))
        if x.is_file()
    ]
    manifest = {
        "experiment": "8-6",
        "track": "orpheus_cross_sentence_voice_consistency",
        "status": "trained_and_generated",
        "seed": SEED,
        "base_model": BASE_MODEL,
        "base_model_revision": HfApi().model_info(BASE_MODEL).sha,
        "dataset": DATASET,
        "dataset_revision": DATASET_REVISION,
        "source_dataset_note": "Public non-disabled mirror of the disabled MrDragonFox/Elise dataset named by the upstream notebook.",
        "train_examples_requested": args.train_examples,
        "train_examples_encoded": len(train_rows),
        "eval_examples_requested": args.eval_examples,
        "eval_examples_encoded": len(eval_rows),
        "train_failures": train_failures,
        "eval_failures": eval_failures,
        "max_audio_seconds": args.max_audio_seconds,
        "optimizer_steps": args.steps,
        "effective_batch_size": 4,
        "lora_rank": 16,
        "pre_eval": pre_eval,
        "train_metrics": train_result.metrics,
        "post_eval": post_eval,
        "gpu": torch.cuda.get_device_name(0),
        "peak_gpu_memory_bytes": torch.cuda.max_memory_reserved(),
        "wall_seconds": time.time() - started,
        "adapter_local_files": adapter_files,
        "adapter_huggingface_repo": f"https://huggingface.co/{args.hf_repo}",
        "adapter_huggingface_revision": adapter_revision,
        "audio": base_audio + adapted_audio,
    }
    (args.output / "orpheus_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()