| |
| """Free-running synthesis gate + ASR scoring for the base model or a LoRA checkpoint. |
| |
| For every sentence x reference (x seed): synthesize with voice cloning (reference in the system |
| prompt, exactly like inference), save the wav, and measure |
| * duration ratio = generated seconds / (words / 2.5) (babble or early stop) |
| * active fraction = share of 100 ms frames with RMS > 0.02 (silence collapse) |
| * tail ratio = RMS(last 25%) / RMS(first 75%) (fade-out collapse, July failure) |
| then transcribe with the Egyptian-dialect Whisper judge and compute normalised WER / CER. |
| |
| Usage (fish-speech root, uv env, HF_HOME=/root/work/hf): |
| python scripts/gate_eval.py --tag base --sentences data/eval/quick.jsonl --out results/eval/base |
| python scripts/gate_eval.py --tag s250 --lora results/<run>/checkpoints/step_000000250.ckpt \ |
| --lora-config r_32_alpha_16_slow --sentences data/eval/quick.jsonl --out results/eval/s250 |
| """ |
| import argparse |
| import json |
| import os |
| import re |
| import sys |
| import time |
| from pathlib import Path |
|
|
| os.environ.setdefault("HF_HOME", "/root/work/hf") |
| os.environ.setdefault("OMP_NUM_THREADS", "8") |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
| import torchaudio |
| from hydra import compose, initialize |
| from hydra.utils import instantiate |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) |
| from fish_speech.models.text2semantic.inference import ( |
| decode_one_token_ar, decode_to_audio, encode_audio, generate_long, load_codec_model, |
| ) |
| from fish_speech.models.text2semantic.llama import BaseTransformer, DualARTransformer |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base", default="checkpoints/s2-pro") |
| ap.add_argument("--lora", default=None, help="Lightning .ckpt with LoRA weights (unmerged eval)") |
| ap.add_argument("--lora-config", default="r_32_alpha_16_slow") |
| ap.add_argument("--merged", default=None, help="merged model dir (alternative to --lora)") |
| ap.add_argument("--tag", required=True) |
| ap.add_argument("--sentences", default="data/eval/quick.jsonl") |
| ap.add_argument("--refs", default="data/eval/refs.jsonl") |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--seeds", type=int, default=1) |
| ap.add_argument("--temperature", type=float, default=0.8) |
| ap.add_argument("--top-p", type=float, default=0.8) |
| ap.add_argument("--max-refs", type=int, default=99) |
| ap.add_argument("--merge-mode", choices=["unmerged", "bf16"], default="unmerged", |
| help="unmerged: keep the LoRA as a separate fp-accurate path (default); bf16: merge into bf16 weights (old behaviour)") |
| ap.add_argument("--no-asr", action="store_true") |
| ap.add_argument("--prefix", default="", help="text prepended to every synthesized sentence and to the reference text (e.g. \"[egyptian accent] \")") |
| ap.add_argument("--judge", choices=["soniox", "whisper"], default=None, |
| help="default: soniox when SONIOX_API_KEY is set, else whisper") |
| ap.add_argument("--asr-model", default="oddadmix/whisper-large-v3-turbo-arabic-dialectal-v2") |
| args = ap.parse_args() |
| judge = args.judge or ("soniox" if os.environ.get("SONIOX_API_KEY") else "whisper") |
| SONIOX_API = os.environ.get("SONIOX_API_BASE", "https://api.soniox.com").rstrip("/") |
| if judge == "soniox": |
| import requests |
| try: |
| code = requests.get(f"{SONIOX_API}/v1/models", timeout=20, |
| headers={"Authorization": f"Bearer {os.environ.get('SONIOX_API_KEY', '')}"}).status_code |
| except Exception as e: |
| code = repr(e) |
| if code != 200: |
| print(f"WARNING: Soniox key check failed ({code}); falling back to the Whisper judge") |
| judge = "whisper" |
|
|
| device = "cuda" |
| out = Path(args.out); (out / "wav").mkdir(parents=True, exist_ok=True) |
| sentences = [json.loads(l) for l in Path(args.sentences).read_text(encoding="utf-8").splitlines() if l.strip()] |
| refs = [json.loads(l) for l in Path(args.refs).read_text(encoding="utf-8").splitlines() if l.strip()][: args.max_refs] |
|
|
| |
| t0 = time.time() |
| if args.lora: |
| with initialize(version_base="1.3", config_path="../fish_speech/configs/lora"): |
| lc = instantiate(compose(config_name=args.lora_config)) |
| model = BaseTransformer.from_pretrained(args.base, load_weights=True, max_length=8192, lora_config=lc) |
| model = model.to(device=device, dtype=torch.bfloat16).train() |
| sd = torch.load(args.lora, map_location="cpu", weights_only=False) |
| sd = sd.get("state_dict", sd) |
| sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items() if "lora" in k} |
| missing, unexpected = model.load_state_dict(sd, strict=False) |
| n_loaded = len(sd) |
| assert not unexpected, unexpected[:5] |
| if args.merge_mode == "unmerged": |
| |
| for m in model.modules(): |
| if hasattr(m, "merge_weights"): |
| m.merge_weights = False |
| model.eval() |
| print(f"loaded {n_loaded} LoRA tensors from {args.lora} (merge mode: {args.merge_mode})") |
| else: |
| src = args.merged or args.base |
| model = DualARTransformer.from_pretrained(src, load_weights=True, max_length=8192) |
| model = model.to(device=device, dtype=torch.bfloat16).eval() |
| model.fixed_temperature = torch.tensor(0.7, device=device) |
| model.fixed_top_p = torch.tensor(0.7, device=device) |
| model.fixed_repetition_penalty = torch.tensor(1.5, device=device) |
| model._cache_setup_done = False |
| with torch.device(device): |
| model.setup_caches(max_batch_size=1, max_seq_len=model.config.max_seq_len, dtype=torch.bfloat16) |
| codec = load_codec_model(Path(args.base) / "codec.pth", device) |
| sr = codec.sample_rate |
| print(f"model ready in {time.time()-t0:.0f}s | codec sr {sr}") |
|
|
| ref_codes = {r["name"]: encode_audio(r["wav"], codec, device).cpu() for r in refs} |
| for r in refs: |
| print(f"ref {r['name']}: {tuple(ref_codes[r['name']].shape)} codes") |
|
|
|
|
| def synth(text, ref): |
| chunks = [] |
| gen = generate_long( |
| model=model, device=device, decode_one_token=decode_one_token_ar, text=text, num_samples=1, |
| max_new_tokens=0, top_p=args.top_p, top_k=30, temperature=args.temperature, |
| repetition_penalty=1.1, compile=False, iterative_prompt=True, chunk_length=300, |
| prompt_text=[args.prefix + ref["text"]], prompt_tokens=[ref_codes[ref["name"]]], |
| ) |
| for resp in gen: |
| if resp.action == "sample" and resp.codes is not None: |
| chunks.append(resp.codes) |
| codes = torch.cat(chunks, dim=1) if len(chunks) > 1 else chunks[0] |
| wav = decode_to_audio(codes.to(device), codec).float().cpu().numpy() |
| return codes.shape[1], wav |
|
|
|
|
| def audio_stats(wav, words): |
| dur = len(wav) / sr |
| frame = int(sr * 0.1) |
| frames = [wav[i:i + frame] for i in range(0, max(frame, len(wav) - frame + 1), frame)] |
| rms = np.array([np.sqrt(np.mean(f ** 2)) for f in frames if len(f)]) |
| active = float((rms > 0.02).mean()) if len(rms) else 0.0 |
| k = max(1, int(0.75 * len(rms))) |
| head = float(np.sqrt(np.mean(rms[:k] ** 2))) if k else 0.0 |
| tail = float(np.sqrt(np.mean(rms[k:] ** 2))) if len(rms) > k else 0.0 |
| return {"dur_s": round(dur, 2), "dur_ratio": round(dur / max(0.5, words / 2.5), 2), |
| "active_frac": round(active, 3), "tail_ratio": round(tail / head, 3) if head > 1e-6 else 0.0} |
|
|
|
|
| rows = [] |
| t1 = time.time() |
| for s in sentences: |
| for ref in refs: |
| for seed in range(args.seeds): |
| torch.manual_seed(seed); torch.cuda.manual_seed(seed) |
| name = f"{s['id']}__{ref['name']}__s{seed}" |
| try: |
| n_codes, wav = synth(args.prefix + s["text"], ref) |
| except Exception as e: |
| print("GEN ERROR", name, repr(e)[:200]); rows.append({"name": name, "error": repr(e)[:200]}); continue |
| path = out / "wav" / f"{name}.wav" |
| sf.write(str(path), np.clip(wav, -1.0, 1.0), sr, subtype="PCM_16") |
| st = audio_stats(wav, s.get("words") or len(s["text"].split())) |
| rows.append({"name": name, "id": s["id"], "src": s.get("src"), "ref": ref["name"], "seed": seed, |
| "text": s["text"], "asr_text": s.get("asr_text"), "n_codes": int(n_codes), "wav": str(path), **st}) |
| print(f"{len(rows):4d}/{len(sentences)*len(refs)*args.seeds} {name} {st['dur_s']}s ratio={st['dur_ratio']} act={st['active_frac']} tail={st['tail_ratio']}", flush=True) |
| print(f"synthesis done in {(time.time()-t1)/60:.1f} min") |
|
|
| |
| DIAC = re.compile(r"[ؐ-ًؚ-ٰٟۖ-ۭـ]") |
| PUNCT = re.compile(r"[^\w\s]", re.UNICODE) |
| AR_DIGITS = str.maketrans("٠١٢٣٤٥٦٧٨٩", "0123456789") |
|
|
|
|
| def norm(t): |
| t = DIAC.sub("", t).translate(AR_DIGITS) |
| t = t.replace("أ", "ا").replace("إ", "ا").replace("آ", "ا").replace("ى", "ي").replace("ة", "ه").replace("ؤ", "و").replace("ئ", "ي") |
| t = PUNCT.sub(" ", t) |
| return " ".join(t.split()) |
|
|
|
|
| def edit_distance(a, b): |
| prev = list(range(len(b) + 1)) |
| for i, x in enumerate(a, 1): |
| cur = [i] |
| for j, y in enumerate(b, 1): |
| cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (x != y))) |
| prev = cur |
| return prev[-1] |
|
|
|
|
| def soniox_transcribe(path, retries=3): |
| """Soniox stt-async-v5 (same engine as the July judge and the podcast transcripts).""" |
| import requests |
| api = SONIOX_API |
| hdr = {"Authorization": f"Bearer {os.environ['SONIOX_API_KEY']}"} |
| last = None |
| for attempt in range(retries): |
| try: |
| with open(path, "rb") as f: |
| r = requests.post(f"{api}/v1/files", headers=hdr, files={"file": (Path(path).name, f, "audio/wav")}, timeout=120) |
| r.raise_for_status() |
| file_id = r.json()["id"] |
| try: |
| r = requests.post(f"{api}/v1/transcriptions", headers=hdr, |
| json={"file_id": file_id, "model": "stt-async-v5", "language_hints": ["ar"]}, timeout=30) |
| r.raise_for_status() |
| tid = r.json()["id"] |
| try: |
| for _ in range(180): |
| s = requests.get(f"{api}/v1/transcriptions/{tid}", headers=hdr, timeout=30).json() |
| if s.get("status") == "completed": |
| break |
| if s.get("status") == "error": |
| raise RuntimeError(f"soniox error: {s.get('error_message')}") |
| time.sleep(1) |
| else: |
| raise TimeoutError("soniox poll timeout") |
| return requests.get(f"{api}/v1/transcriptions/{tid}/transcript", headers=hdr, timeout=30).json().get("text", "") |
| finally: |
| requests.delete(f"{api}/v1/transcriptions/{tid}", headers=hdr, timeout=30) |
| finally: |
| requests.delete(f"{api}/v1/files/{file_id}", headers=hdr, timeout=30) |
| except Exception as e: |
| last = e |
| time.sleep(3 * (attempt + 1)) |
| raise RuntimeError(f"soniox failed after {retries} attempts: {last!r}") |
|
|
|
|
| def score(r, hyp): |
| if r.get("asr_text"): |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| from egy_numbers import verbalize |
| ref_n, hyp_n = norm(verbalize(r["asr_text"])), norm(verbalize(hyp)) |
| else: |
| ref_n, hyp_n = norm(r["text"]), norm(hyp) |
| rw, hw = ref_n.split(), hyp_n.split() |
| r["hyp"] = hyp |
| r["wer"] = round(edit_distance(rw, hw) / max(1, len(rw)), 4) |
| r["cer"] = round(edit_distance(ref_n.replace(" ", ""), hyp_n.replace(" ", "")) / max(1, len(ref_n.replace(" ", ""))), 4) |
|
|
|
|
| if not args.no_asr and rows: |
| todo = [r for r in rows if "wav" in r] |
| print(f"ASR judge: {judge} on {len(todo)} files") |
| if judge == "soniox": |
| from concurrent.futures import ThreadPoolExecutor |
|
|
| def one(r): |
| try: |
| return soniox_transcribe(r["wav"]) |
| except Exception as e: |
| return e |
|
|
| with ThreadPoolExecutor(max_workers=8) as ex: |
| for r, hyp in zip(todo, ex.map(one, todo)): |
| if isinstance(hyp, Exception): |
| r["asr_error"] = repr(hyp)[:200]; print("ASR ERROR", r["name"], r["asr_error"]) |
| else: |
| score(r, hyp) |
| else: |
| del model; torch.cuda.empty_cache() |
| from transformers import pipeline |
| asr = pipeline("automatic-speech-recognition", model=args.asr_model, torch_dtype=torch.float16, device=0) |
| for r in todo: |
| wav, s_r = sf.read(r["wav"], dtype="float32") |
| wav16 = torchaudio.functional.resample(torch.from_numpy(wav), s_r, 16000).numpy() |
| hyp = asr({"raw": wav16, "sampling_rate": 16000}, chunk_length_s=30, |
| generate_kwargs={"language": "ar", "task": "transcribe"})["text"] |
| score(r, hyp) |
|
|
| |
| ok = [r for r in rows if "wav" in r] |
| scored = [r for r in ok if "wer" in r] |
| def mean(k, rs=None): rs = ok if rs is None else rs; return round(float(np.mean([r[k] for r in rs])), 4) if rs else None |
| summary = { |
| "tag": args.tag, "judge": judge if not args.no_asr else None, "n": len(ok), "errors": len(rows) - len(ok), |
| "n_scored": len(scored), "wer": mean("wer", scored), "cer": mean("cer", scored), |
| "dur_ratio_median": round(float(np.median([r["dur_ratio"] for r in ok])), 3) if ok else None, |
| "active_frac_mean": mean("active_frac"), "tail_ratio_mean": mean("tail_ratio"), |
| "collapsed": sum(1 for r in ok if r["tail_ratio"] < 0.3 or r["active_frac"] < 0.5), |
| "too_long": sum(1 for r in ok if r["dur_ratio"] > 1.6), "too_short": sum(1 for r in ok if r["dur_ratio"] < 0.5), |
| } |
| for src in sorted({r["src"] for r in scored}): |
| sub = [r for r in scored if r["src"] == src] |
| summary[f"wer_{src}"] = round(float(np.mean([r["wer"] for r in sub])), 4) |
| for ref in sorted({r["ref"] for r in scored}): |
| sub = [r for r in scored if r["ref"] == ref] |
| summary[f"wer_ref_{ref}"] = round(float(np.mean([r["wer"] for r in sub])), 4) |
| (out / "rows.jsonl").write_text("\n".join(json.dumps(r, ensure_ascii=False) for r in rows) + "\n", encoding="utf-8") |
| (out / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") |
| print("SUMMARY", json.dumps(summary, ensure_ascii=False)) |
|
|
