Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
664 lines
39 KiB
Python
664 lines
39 KiB
Python
#!/usr/bin/env python3
|
||
"""Golden numerical fixtures for the lfm3_audio C++ port (LFM3-Audio-2B, HF remote code).
|
||
|
||
make_fixtures.py --model <model dir> --assets <dir with c.wav, speech.wav, a.wav> --out <fixture dir>
|
||
|
||
Reference = HF model in float32 on CUDA (sdpa, TF32 off). Key stages are repeated with a bf16 copy of the model
|
||
to measure the bf16-vs-fp32 noise floor (manifest["noise_floor"]).
|
||
Every array is a .npy file listed in <out>/manifest.json with shape, dtype, meaning and inputs.
|
||
The fixtures are derived from the model weights: keep them out of git.
|
||
"""
|
||
import argparse
|
||
import contextlib
|
||
import json
|
||
import os
|
||
import time
|
||
|
||
import librosa
|
||
import numpy as np
|
||
import soundfile as sf
|
||
import torch
|
||
from transformers import AutoModelForMultimodalLM, AutoProcessor, DynamicCache
|
||
from transformers.dynamic_module_utils import get_class_from_dynamic_module
|
||
|
||
_args = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
_args.add_argument("--model", required=True)
|
||
_args.add_argument("--assets", required=True)
|
||
_args.add_argument("--out", required=True)
|
||
ARGS = _args.parse_args()
|
||
MID, ASSETS, OUT = ARGS.model, ARGS.assets, ARGS.out
|
||
DEV = "cuda"
|
||
ORDER_SEED = 20261024 # numpy default_rng seed of the injected MaskGIT orders (reset per run)
|
||
S2S_ORDER_SEED = 3 # with this seed greedy S2S ends with an EOS frame + <|im_end|>; ORDER_SEED never emits EOS
|
||
MAX_LOGIT_FRAMES = 40 # depth logits are stored for the first frames of each generation only
|
||
# generate() takes the speech sampling it is not given from generation_config.json, which samples, so greedy
|
||
# speech is asked for.
|
||
GREEDY = {"audio_temperature": 0.0, "audio_repetition_penalty": 1.0}
|
||
|
||
torch.backends.cuda.matmul.allow_tf32 = False
|
||
torch.backends.cudnn.allow_tf32 = False
|
||
torch.set_float32_matmul_precision("highest")
|
||
|
||
AUDIO_TOKEN, AUDIO_OUT_TOKEN, EOS_CODE = 125031, 125032, 2016
|
||
IM_END, NEWLINE = 124900, 207
|
||
S2S_SYSTEM = "Respond with interleaved text and audio."
|
||
M = {"meta": {}, "inputs": {}, "prompts": {}, "arrays": {}, "noise_floor": {}, "results": {}}
|
||
|
||
|
||
def np_of(x):
|
||
if torch.is_tensor(x):
|
||
x = x.detach().cpu()
|
||
x = x.float() if x.is_floating_point() else x
|
||
x = x.numpy()
|
||
x = np.asarray(x)
|
||
if x.dtype.kind in "iub":
|
||
return x.astype(np.int32)
|
||
return x.astype(np.float32) if x.dtype != np.float32 else x
|
||
|
||
|
||
def save(name, x, meaning, **inputs):
|
||
a = np.ascontiguousarray(np_of(x))
|
||
path = os.path.join(OUT, name + ".npy")
|
||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||
np.save(path, a)
|
||
M["arrays"][name] = {"file": name + ".npy", "shape": list(a.shape), "dtype": str(a.dtype), "meaning": meaning,
|
||
"inputs": inputs}
|
||
return a
|
||
|
||
|
||
def rel_rms(a, b):
|
||
a, b = np_of(a).astype(np.float64), np_of(b).astype(np.float64)
|
||
return float(np.sqrt(np.mean((a - b) ** 2)) / max(np.sqrt(np.mean(b ** 2)), 1e-30))
|
||
|
||
|
||
def floor(stage, bf, fp, **extra):
|
||
bf_, fp_ = np_of(bf).astype(np.float64), np_of(fp).astype(np.float64)
|
||
d = {"rel_rms": rel_rms(bf_, fp_), "max_abs": float(np.abs(bf_ - fp_).max()), "ref_rms": float(np.sqrt(np.mean(fp_ ** 2)))}
|
||
d.update(extra)
|
||
M["noise_floor"][stage] = d
|
||
print(f" floor {stage}: {d}", flush=True)
|
||
|
||
|
||
def first(o):
|
||
return o[0] if isinstance(o, (tuple, list)) else o
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def capture(mods):
|
||
"""Records the first output of each named module."""
|
||
got = {}
|
||
hooks = [m.register_forward_hook(lambda mod, a, o, k=k: got.setdefault(k, first(o).detach().clone())) for k, m in mods.items()]
|
||
try:
|
||
yield got
|
||
finally:
|
||
for h in hooks:
|
||
h.remove()
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def frame_recorder(model, inject_seed):
|
||
"""Wraps sample_audio_frame: injects the MaskGIT order and records the depth decoder passes."""
|
||
rec = {"hidden": [], "order": [], "codes": [], "pass_codes": [], "pass_masked": [], "pass_logits": []}
|
||
rng = np.random.default_rng(inject_seed)
|
||
orig_sample, orig_randperm = model.sample_audio_frame, torch.randperm
|
||
passes = []
|
||
|
||
def on_depth(mod, args, out):
|
||
_, codes, masked = args
|
||
passes.append((codes[0].detach().clone(), masked[0].detach().clone(), out[0].detach().float().clone()))
|
||
|
||
def sample(hidden_state, *a, **kw):
|
||
assert hidden_state.ndim == 1, "batch 1 only"
|
||
order = torch.as_tensor(rng.permutation(12), dtype=torch.long)
|
||
torch.randperm = lambda n, device=None, **_: order.to(device)
|
||
passes.clear()
|
||
try:
|
||
codes = orig_sample(hidden_state, *a, **kw)
|
||
finally:
|
||
torch.randperm = orig_randperm
|
||
rec["order"].append(order.clone())
|
||
assert len(passes) == 3
|
||
rec["hidden"].append(hidden_state.detach().float().cpu().clone())
|
||
rec["codes"].append(codes.detach().cpu().clone())
|
||
rec["pass_codes"].append(torch.stack([p[0] for p in passes]).cpu())
|
||
rec["pass_masked"].append(torch.stack([p[1] for p in passes]).cpu())
|
||
rec["pass_logits"].append(torch.stack([p[2] for p in passes]).cpu())
|
||
return codes
|
||
|
||
h = model.model.depth_decoder.register_forward_hook(on_depth)
|
||
model.sample_audio_frame = sample
|
||
try:
|
||
yield rec
|
||
finally:
|
||
h.remove()
|
||
del model.sample_audio_frame # back to the class method
|
||
torch.randperm = orig_randperm
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def text_logit_recorder(model, k=20):
|
||
rec = []
|
||
h = model.lm_head.register_forward_hook(
|
||
lambda mod, a, o: rec.append(torch.topk(o.detach().float().reshape(-1), k)) if o.ndim == 1 else None)
|
||
try:
|
||
yield rec
|
||
finally:
|
||
h.remove()
|
||
|
||
|
||
def slots_of(out, prompt_len):
|
||
"""Per generated slot: kind (0 text, 1 audio), token id (125032 for frames) and codes (-1 for text)."""
|
||
seq = out.sequences[0, prompt_len:].tolist()
|
||
codes = out.audio_codes[0].cpu()
|
||
kinds, toks, frames, j = [], [], [], 0
|
||
for t in seq:
|
||
kinds.append(int(t == AUDIO_OUT_TOKEN))
|
||
toks.append(t)
|
||
if t == AUDIO_OUT_TOKEN:
|
||
frames.append(codes[:, j].tolist())
|
||
j += 1
|
||
else:
|
||
frames.append([-1] * 12)
|
||
assert j == int(out.audio_lengths[0])
|
||
pattern, i = [], 0
|
||
while i < len(seq):
|
||
e = i
|
||
while e < len(seq) and kinds[e] == kinds[i]:
|
||
e += 1
|
||
pattern.append(f"A{e - i}" if kinds[i] else f"T{e - i}")
|
||
i = e
|
||
return np.array(kinds), np.array(toks), np.array(frames), " ".join(pattern)
|
||
|
||
|
||
def multiturn(model, proc):
|
||
"""Turn 2 of an S2S conversation (user audio a.wav) after turn 1 = the gen/s2s_c_greedy(_trunc) slots.
|
||
|
||
The fixture is the re-prefill of the replayed history: turn-1 prompt + turn-1 slots exactly as
|
||
generated (+ an EOS frame when the cut turn ends without one, + <|im_end|> when the turn was cut) + "\\n" + the
|
||
turn-2 user block and <|mixed_start|>. It is checked against HF cache continuation (turn 1 regenerated with
|
||
past_key_values, then only the suffix fed) and against the HF chat template rendering the same history."""
|
||
tok = proc.tokenizer
|
||
a = librosa.load(f"{ASSETS}/a.wav", sr=16000, mono=True)[0]
|
||
sf.write(f"{OUT}/audio/a_16k.wav", a, 16000, subtype="FLOAT")
|
||
assert np.array_equal(sf.read(f"{OUT}/audio/a_16k.wav", dtype="float32")[0], a)
|
||
M["inputs"]["a"] = {"file": "audio/a_16k.wav", "samples": int(len(a)), "sample_rate": 16000,
|
||
"source": f"{ASSETS}/a.wav librosa.load(sr=16000, mono=True, soxr_hq)"}
|
||
c_path = f"{OUT}/audio/c_16k.wav"
|
||
c = sf.read(c_path, dtype="float32")[0]
|
||
conv1 = [{"role": "system", "content": S2S_SYSTEM}, {"role": "user", "content": [{"type": "audio", "audio": c_path}]},
|
||
{"role": "assistant", "content": "<|mixed_start|>"}]
|
||
inp1 = proc.apply_chat_template(conv1, continue_final_message=True, tokenize=True, return_dict=True,
|
||
return_tensors="pt").to(DEV)
|
||
ids1 = inp1["input_ids"][0].tolist()
|
||
assert ids1 == np.load(f"{OUT}/prompt/s2s_c_ids.npy").tolist(), "turn-1 prompt differs from prompt/s2s_c_ids"
|
||
tail = [{"role": "user", "content": [{"type": "audio", "audio": f"{OUT}/audio/a_16k.wav"}]},
|
||
{"role": "assistant", "content": "<|mixed_start|>"}]
|
||
# the template's "<|im_end|>\n" after turn 1, then the turn-2 blocks (an empty system block renders nothing)
|
||
suffix_text = "\n" + proc.apply_chat_template([{"role": "system", "content": ""}] + tail, continue_final_message=True,
|
||
tokenize=False)[len(tok.bos_token):]
|
||
assert suffix_text.startswith("\n<|im_start|>user\n"), suffix_text[:40]
|
||
|
||
def encode(text, clips):
|
||
b = proc(text=text, audio=clips, return_tensors="pt")
|
||
return {k: b[k].to(DEV) for k in ("input_ids", "input_features", "feature_attention_mask") if k in b}
|
||
|
||
def fwd(ids, feats, codes, cache):
|
||
past = cache.get_seq_length()
|
||
L = ids.shape[1]
|
||
with torch.no_grad():
|
||
o = model.model(input_ids=ids, input_features=feats.get("input_features"),
|
||
feature_attention_mask=feats.get("feature_attention_mask"), audio_out_codes=codes,
|
||
attention_mask=torch.ones(1, past + L, dtype=torch.long, device=DEV),
|
||
position_ids=torch.arange(past, past + L, device=DEV)[None], past_key_values=cache, use_cache=True)
|
||
return o.last_hidden_state[0, -1].float(), model.lm_head(o.last_hidden_state[:, -1])[0].float()
|
||
|
||
suf = encode(suffix_text[1:], [a])["input_ids"][0].tolist() # the "\n" is appended with [NEWLINE]
|
||
both = encode("<|audio|><|audio|>", [c, a]) # features of both clips, batched as the processor does
|
||
for tag, gen_tag, order_seed, max_new in (("multiturn_s2s_turn2", "s2s_c_greedy", S2S_ORDER_SEED, 1024),
|
||
("multiturn_s2s_trunc_turn2", "s2s_c_greedy_trunc", ORDER_SEED, 96)):
|
||
gen = np.load(f"{OUT}/gen/{gen_tag}_slot_token.npy").tolist()
|
||
kinds = np.load(f"{OUT}/gen/{gen_tag}_slot_kind.npy")
|
||
frames = np.load(f"{OUT}/gen/{gen_tag}_slot_codes.npy")[kinds == 1].tolist()
|
||
appended = []
|
||
if gen[-1] != IM_END:
|
||
if frames and frames[-1][0] != EOS_CODE:
|
||
appended.append(AUDIO_OUT_TOKEN)
|
||
frames.append([EOS_CODE] * 12)
|
||
appended.append(IM_END)
|
||
ids = ids1 + gen + appended + [NEWLINE] + suf
|
||
codes = torch.tensor(frames, device=DEV)
|
||
# (b) re-prefill of the whole conversation: the fixture
|
||
hb, lb = fwd(torch.tensor([ids], device=DEV), both, codes, DynamicCache(config=model.config.text_config))
|
||
# (a) cache continuation: regenerate turn 1 with a caller cache, then feed only what follows the slots
|
||
cache = DynamicCache(config=model.config.text_config)
|
||
with frame_recorder(model, order_seed), torch.no_grad():
|
||
o = model.generate(**inp1, past_key_values=cache, return_audio_eos=True, max_new_tokens=max_new, **GREEDY)
|
||
_, toks, fr, _ = slots_of(o, len(ids1))
|
||
regenerated_equal = toks.tolist() == gen and np.array_equal(fr, np.load(f"{OUT}/gen/{gen_tag}_slot_codes.npy"))
|
||
cache_len = cache.get_seq_length()
|
||
sa = torch.tensor([appended + [NEWLINE] + suf], device=DEV)
|
||
a_feats = encode(suffix_text[1:], [a])
|
||
ha, la = fwd(sa, a_feats, codes[-1:] if AUDIO_OUT_TOKEN in appended else None, cache)
|
||
# (c) chat-template re-render: assistant content = prefill + decoded slots (+ the EOS frame placeholder)
|
||
body = gen[:-1] if gen[-1] == IM_END else gen + appended[:-1]
|
||
content = "<|mixed_start|>" + tok.decode(body, skip_special_tokens=False)
|
||
conv = conv1[:2] + [{"role": "assistant", "content": content}] + tail
|
||
text = proc.apply_chat_template(conv, continue_final_message=True, tokenize=False)
|
||
ids_tpl = encode(text, [c, a])["input_ids"][0].tolist()
|
||
top = torch.topk(lb, 20)
|
||
save(f"gen/{tag}_prompt_ids", np.array(ids), "turn-2 prompt: turn-1 prompt + turn-1 slots (+ EOS frame + <|im_end|> "
|
||
"if cut) + \\n + user a.wav + <|mixed_start|>", turn1=f"gen/{gen_tag}_slot_*", audio=["audio/c_16k.wav", "audio/a_16k.wav"])
|
||
save(f"gen/{tag}_last_logits", lb, "last-position logits of the turn-2 prompt, re-prefilled [128000]",
|
||
prompt=f"gen/{tag}_prompt_ids")
|
||
save(f"gen/{tag}_top20_ids", top.indices, "top-20 ids of last_logits", prompt=f"gen/{tag}_prompt_ids")
|
||
r = {"turn1": gen_tag, "prompt_len": len(ids), "appended_after_turn1": appended + [NEWLINE],
|
||
"turn1_regenerated_with_cache_equal_fixture": bool(regenerated_equal), "cache_len_after_turn1": int(cache_len),
|
||
"cache_covers_prompt_plus_slots": cache_len == len(ids1) + len(gen),
|
||
"cache_continuation_vs_reprefill": {"logits_max_abs": float((la - lb).abs().max()),
|
||
"hidden_max_abs": float((ha - hb).abs().max()),
|
||
"hidden_rel": float((ha - hb).norm() / hb.norm()),
|
||
"argmax_same": int(la.argmax()) == int(lb.argmax()),
|
||
"top20_same": torch.topk(la, 20).indices.tolist() == top.indices.tolist()},
|
||
"template_render_ids_equal": ids_tpl == ids}
|
||
if ids_tpl != ids:
|
||
i = next((k for k in range(min(len(ids), len(ids_tpl))) if ids[k] != ids_tpl[k]), min(len(ids), len(ids_tpl)))
|
||
r["template_first_diff"] = {"pos": i, "template": ids_tpl[i:i + 8], "replay": ids[i:i + 8]}
|
||
M["results"][tag] = r
|
||
print(f" {tag}: {r}", flush=True)
|
||
|
||
|
||
def dump_probe(name, obj, indent=1):
|
||
os.makedirs(f"{OUT}/probes", exist_ok=True)
|
||
with open(f"{OUT}/probes/{name}", "w") as f:
|
||
json.dump(obj, f, ensure_ascii=False, indent=indent)
|
||
|
||
|
||
def tokenizer_cases(proc):
|
||
"""probes/tokenizer_cases.json: ids, tokens and decodes of HF's tokenizer (as the processor loads it) for a regression set."""
|
||
tok = proc.tokenizer
|
||
added = sorted(tok.added_tokens_decoder.items()) # (id, AddedToken)
|
||
lang = {
|
||
"arabic": ["مرحبا، كيف حالك اليوم؟", "القطار التالي يغادر خلال عشر دقائق.", "أنا أحب القراءة ١٢٣٤٥"],
|
||
"danish": ["Hej, hvordan går det i dag?", "Rødgrød med fløde på ærøskøbing.", "Næste tog kører om ti minutter."],
|
||
"german": ["Guten Tag, wie geht's?", "Die Straße ist größer als gedacht.", "Donaudampfschifffahrtsgesellschaftskapitän"],
|
||
"greek": ["Γειά σου κόσμε!", "Το επόμενο τρένο φεύγει σε δέκα λεπτά.", "Ψυχή, ἀρετή, ὕδωρ"],
|
||
"british_english": ["Colour, favour and organisation.", "I'll queue at the centre, mate.", "WHAT'S HE DOING? THEY'LL SEE."],
|
||
"american_english": ["The quick brown fox jumps over the lazy dog", "I'm sure you've heard: it's 3:45 p.m.", "Don't, won't, can't, shouldn't've"],
|
||
"finnish": ["Hyvää huomenta, mitä kuuluu?", "Lentokonesuihkuturbiinimoottoriapumekaanikkoaliupseerioppilas", "Seuraava juna lähtee kymmenen minuutin päästä."],
|
||
"french": ["Bonjour, ça va ? C'est l'été !", "L'œuvre d'art coûte 1 234,56 €.", "Qu'est-ce que c'est ?"],
|
||
"hebrew": ["שלום, מה שלומך?", "הרכבת הבאה יוצאת בעוד עשר דקות.", "בְּרֵאשִׁית בָּרָא אֱלֹהִים"],
|
||
"italian": ["Buongiorno, come stai?", "Perché l'università è chiusa?", "Il prossimo treno parte tra dieci minuti."],
|
||
"japanese": ["こんにちは、元気ですか？", "次の電車は十分後に出発します。", "カタカナとひらがなと漢字、ｶﾀｶﾅ半角", "東京タワーは333メートルです。"],
|
||
"korean": ["안녕하세요, 잘 지내세요?", "다음 기차는 10분 후에 출발합니다.", "한국어 자모: ㄱㄴㄷ 가"],
|
||
"dutch": ["Goedemorgen, hoe gaat het?", "De volgende trein vertrekt over tien minuten.", "'s Ochtends en 't is koud."],
|
||
"polish": ["Dzień dobry, jak się masz?", "Zażółć gęślą jaźń.", "Następny pociąg odjeżdża za dziesięć minut."],
|
||
"portuguese": ["Olá, tudo bem? Não sei.", "O próximo comboio parte em dez minutos.", "Ação, coração, pão."],
|
||
"spanish": ["¡Hola! ¿Cómo estás?", "El niño pequeño comió piña.", "El próximo tren sale en diez minutos."],
|
||
"swedish": ["Hej, hur mår du?", "Nästa tåg avgår om tio minuter.", "Smörgåsbord och blåbär."],
|
||
"turkish": ["Merhaba, nasılsın?", "İstanbul'da ılık bir gün. IĞDIR ığdır", "Bir sonraki tren on dakika içinde kalkıyor."],
|
||
"chinese": ["你好，世界！", "下一班火车十分钟后出发。", "繁體中文與简体中文混合", "数字12345 can't"],
|
||
}
|
||
cases = []
|
||
def add(cat, s):
|
||
cases.append({"category": cat, "text": s})
|
||
|
||
for l, ss in lang.items():
|
||
for s in ss:
|
||
add("lang:" + l, s)
|
||
|
||
for s in ["1", "12", "123", "1234", "12345", "1234567890", "3.14159", "1,000,000", "2026-10-24", "+1 (555) 010-9999",
|
||
"٣٤٥٦", "１２３４５", "a1b22c333d4444", "007", "1st 2nd 3rd 4th", "0x1F600", "100%", "$5.99", "-273.15°C"]:
|
||
add("digits", s)
|
||
|
||
for s in ["IT'S", "it's", "It'S", "HE'LL", "he'll", "THEY'RE", "we've", "WE'VE", "I'M", "i'd", "I'D", "don't", "DON'T",
|
||
"'s", "'S", "'ll", "'LL", "o'clock", "rock'n'roll", "The cat's toy's colour", "l’homme", "it’s (curly)", "''quoted''"]:
|
||
add("contractions", s)
|
||
|
||
for s in ["a\r\nb", "a\r\n\r\nb", "\r\n", "\r\n\r\n\r\n", "line1\nline2\n", "\n\n\n", "a\rb", "x \r\n y", "end.\r\n",
|
||
" \n \n", "\n", " \n", "\n "]:
|
||
add("newlines", s)
|
||
|
||
for s in [" leading", "trailing ", " two leading", "two trailing ", " ", " ", " ", "a b", "a b", "a b",
|
||
" Hello world\n\nfoo", "\u00a0nbsp\u00a0", "ideographic\u3000space", "zero\u200bwidth"]:
|
||
add("spaces", s)
|
||
|
||
for s in ["\t", "a\tb", "\t\tindented", "col1\tcol2\tcol3", "\t \n\t", "def f():\n\treturn 1"]:
|
||
add("tabs", s)
|
||
|
||
for s in ["😀", "👍🏽", "👨\u200d👩\u200d👧\u200d👦", "🇯🇵🇺🇸", "I ❤\ufe0f NY", "emoji😀in😀text", "✓ ✗ ★ ☆", "🙂🙂🙂", "©®™", "→←↑↓"]:
|
||
add("emoji", s)
|
||
|
||
for s in ["e\u0301", "café vs cafe\u0301", "n\u0303", "a\u0300\u0301\u0302", "Z\u0351\u0352\u0353algo", "क\u094dष",
|
||
"fi ligature", "Å vs A\u030a", "Dž title"]:
|
||
add("combining", s)
|
||
|
||
for s in ["<function=x>", "<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>", "<think>", "</think>",
|
||
"<think>I should answer.</think>The answer is 4.", "<think></think>", "<tool_call>\n<function=f>\n</function>\n</tool_call>",
|
||
"<|im_start|>user\nhi<|im_end|>\n", "<|im_end|><|im_end|>", "text<|text_end|>", "<|audio|><|audio|>",
|
||
"<|audio_out|>x", "a<|audio_start|>b", "<|mixed_start|>Hello<|text_end|>", "<|startoftext|>", "<|im_start|",
|
||
"<|im_end", "< |im_end|>", "<|IM_END|>", "<think >", "<thinking>", "</think >", "<tools>[]</tools>", "<|pad|>"]:
|
||
add("special_in_text", s)
|
||
|
||
for _, t in added:
|
||
add("added_token_embedded", f"pre {t.content} post")
|
||
add("added_token_bare", t.content)
|
||
|
||
for s in ["Pneumonoultramicroscopicsilicovolcanoconiosis", "a" * 100, "ab" * 60, "supercalifragilisticexpialidocious!",
|
||
"https://example.com/path?query=1&b=2#frag", "user@example.com", "snake_case_identifier_name", "CamelCaseIdentifierName",
|
||
"x" * 300, "ﷺ", "𝔘𝔫𝔦𝔠𝔬𝔡𝔢", "\u0000nul", "\x7fdel", "\ufeffbom", "\U0010ffff"]:
|
||
add("long_or_odd", s)
|
||
|
||
for s in ["", "Hello.", "Hello there.", "Perform TTS in American English using a female voice.", "Perform ASR in English.",
|
||
"Respond with interleaved text and audio.", "Perform TTS in Japanses using a male voice.", "Read this with a female voice in Korean."]:
|
||
add("prompts", s)
|
||
|
||
out = []
|
||
for c in cases:
|
||
ids = tok(c["text"], add_special_tokens=False)["input_ids"]
|
||
ids_sp = tok(c["text"], add_special_tokens=True)["input_ids"]
|
||
out.append({**c, "ids": ids,
|
||
"decode_keep_special": tok.decode(ids, skip_special_tokens=False),
|
||
"decode_skip_special": tok.decode(ids, skip_special_tokens=True),
|
||
"tokens": tok.convert_ids_to_tokens(ids),
|
||
"add_special_tokens_changes_ids": ids_sp != ids})
|
||
meta = {"tokenizer_class": type(tok).__name__, "len_tokenizer": len(tok), "n_added": len(added),
|
||
"added_tokens": [{"id": i, "content": t.content, "special": t.special} for i, t in added],
|
||
"loaded_via": "AutoProcessor.from_pretrained(model, trust_remote_code=True).tokenizer; tok(text, add_special_tokens=False)",
|
||
"n_cases": len(out)}
|
||
dump_probe("tokenizer_cases.json", {"meta": meta, "cases": out})
|
||
print(f" tokenizer_cases: {len(out)} cases", flush=True)
|
||
|
||
|
||
def ignore_merges_cases():
|
||
"""probes/ignore_merges_cases.json: every base-vocab string whose ids differ between tokenizer.json's
|
||
ignore_merges=true (as shipped) and false. Ids and tokens are the shipped ones."""
|
||
from tokenizers import Tokenizer
|
||
with open(f"{MID}/tokenizer.json") as f:
|
||
j = json.load(f)
|
||
shipped = Tokenizer.from_str(json.dumps(j))
|
||
j["model"]["ignore_merges"] = False
|
||
merges_only = Tokenizer.from_str(json.dumps(j))
|
||
n = len(j["model"]["vocab"])
|
||
cases = []
|
||
for i in range(n):
|
||
t = shipped.decode([i])
|
||
ea, eb = shipped.encode(t, add_special_tokens=False), merges_only.encode(t, add_special_tokens=False)
|
||
if ea.ids != eb.ids:
|
||
cases.append({"text": t, "ids": ea.ids, "tokens": ea.tokens, "ids_without_ignore_merges": eb.ids})
|
||
dump_probe("ignore_merges_cases.json", {"meta": {"vocab": n, "n_cases": len(cases), "source": "make_fixtures.py"}, "cases": cases}, None)
|
||
print(f" ignore_merges_cases: {len(cases)} of {n}", flush=True)
|
||
|
||
|
||
def main():
|
||
t0 = time.time()
|
||
os.makedirs(OUT, exist_ok=True)
|
||
M["meta"] = {"model": MID, "torch": torch.__version__, "transformers": __import__("transformers").__version__,
|
||
"librosa": librosa.__version__, "device": torch.cuda.get_device_name(0),
|
||
"reference": "fp32 weights (upcast from bf16 checkpoint), attn sdpa, TF32 off",
|
||
"order_seed": ORDER_SEED, "max_logit_frames": MAX_LOGIT_FRAMES,
|
||
"rel_rms": "sqrt(mean((bf16-fp32)^2))/sqrt(mean(fp32^2))"}
|
||
|
||
wavs = {}
|
||
for name in ("c", "speech"):
|
||
y, _ = librosa.load(f"{ASSETS}/{name}.wav", sr=16000, mono=True) # res_type soxr_hq, like the HF processor
|
||
wavs[name] = y
|
||
sr = 16000
|
||
rs = np.random.default_rng(7)
|
||
for name, n in (("syn_1s", 16000), ("syn_037s", 5921)):
|
||
t = np.arange(n) / sr
|
||
y = 0.3 * np.sin(2 * np.pi * (200 * t + 900 * t ** 2)) + 0.05 * rs.standard_normal(n)
|
||
wavs[name] = y.astype(np.float32)
|
||
os.makedirs(f"{OUT}/audio", exist_ok=True)
|
||
for name, y in wavs.items():
|
||
p = f"{OUT}/audio/{name}_16k.wav"
|
||
sf.write(p, y, sr, subtype="FLOAT")
|
||
back, _ = sf.read(p, dtype="float32")
|
||
assert np.array_equal(back, y)
|
||
M["inputs"][name] = {"file": f"audio/{name}_16k.wav", "samples": int(len(y)), "sample_rate": sr,
|
||
"source": f"{ASSETS}/{name}.wav librosa.load(sr=16000, mono=True, soxr_hq)"
|
||
if not name.startswith("syn") else
|
||
"0.3*sin(2pi(200t+900t^2)) + 0.05*N(0,1) (numpy default_rng(7), syn_1s first)"}
|
||
|
||
proc = AutoProcessor.from_pretrained(MID, trust_remote_code=True)
|
||
fe = proc.feature_extractor
|
||
feats = {}
|
||
for name, y in wavs.items():
|
||
b = fe(y, sampling_rate=16000)
|
||
f = b["input_features"][0]
|
||
feats[name] = f
|
||
save(f"feat/{name}_mel", f, "normalized log-mel features [128, T] = (clamp(log10(mel),-12)+4)/4", wav=f"audio/{name}_16k.wav")
|
||
|
||
print("loading fp32 model", flush=True)
|
||
model = AutoModelForMultimodalLM.from_pretrained(MID, dtype=torch.float32, device_map=DEV, trust_remote_code=True,
|
||
attn_implementation="sdpa").eval()
|
||
Codec = get_class_from_dynamic_module("modeling_liquid_codec.LiquidCodecModel", MID)
|
||
SCache = get_class_from_dynamic_module("modeling_liquid_codec.LiquidCodecStreamingCache", MID)
|
||
codec = Codec.from_pretrained(MID, subfolder="audio_tokenizer", dtype=torch.float32, device_map=DEV,
|
||
attn_implementation="sdpa").eval()
|
||
print("loading bf16 model", flush=True)
|
||
model_bf = AutoModelForMultimodalLM.from_pretrained(MID, dtype=torch.bfloat16, device_map=DEV, trust_remote_code=True,
|
||
attn_implementation="sdpa").eval()
|
||
codec_bf = Codec.from_pretrained(MID, subfolder="audio_tokenizer", dtype=torch.bfloat16, device_map=DEV,
|
||
attn_implementation="sdpa").eval()
|
||
M["meta"]["attn_implementation"] = model.config._attn_implementation
|
||
|
||
def enc_mods(m):
|
||
e = m.model.audio_encoder
|
||
return {"downsample": e.downsample, "final_norm": e.lfm.embedding_norm, "projector": m.model.audio_projector}
|
||
|
||
print("encoder", flush=True)
|
||
with torch.no_grad():
|
||
for name, f in feats.items():
|
||
ft = torch.as_tensor(f, device=DEV)[None]
|
||
mask = torch.ones(1, ft.shape[-1], dtype=torch.long, device=DEV)
|
||
with capture(enc_mods(model)) as g:
|
||
model.model.get_audio_features(ft, mask)
|
||
src = {"features": f"feat/{name}_mel"}
|
||
save(f"enc/{name}_downsample", g["downsample"][0].T, "downsampler output (after lin_out) [N, 768]; "
|
||
"flatten index c*16+f before lin_out", **src)
|
||
save(f"enc/{name}_final_norm", g["final_norm"][0], "encoder output after embedding_norm [N, 768]", **src)
|
||
save(f"enc/{name}_projector", g["projector"], "projector output [N, 2048] (scattered into <|audio|> slots)", **src)
|
||
if name == "c":
|
||
with capture(enc_mods(model_bf)) as gb:
|
||
model_bf.model.get_audio_features(ft, mask)
|
||
for k in g:
|
||
floor(f"encoder.{k}", gb[k], g[k])
|
||
|
||
A = {n: f"{OUT}/audio/{n}_16k.wav" for n in wavs}
|
||
convs = {
|
||
"asr_en_c": ([{"role": "system", "content": "Perform ASR in English."},
|
||
{"role": "user", "content": [{"type": "audio", "audio": A["c"]}]}], False),
|
||
"asr_plain_c": ([{"role": "system", "content": "Perform ASR."},
|
||
{"role": "user", "content": [{"type": "audio", "audio": A["c"]}]}], False),
|
||
"asr_plain_speech": ([{"role": "system", "content": "Perform ASR."},
|
||
{"role": "user", "content": [{"type": "audio", "audio": A["speech"]}]}], False),
|
||
"tts_fox": ([{"role": "system", "content": "Perform TTS in American English using a female voice."},
|
||
{"role": "user", "content": "The quick brown fox jumps over the lazy dog"},
|
||
{"role": "assistant", "content": "<|audio_start|>"}], True),
|
||
"s2s_c": ([{"role": "system", "content": "Respond with interleaved text and audio."},
|
||
{"role": "user", "content": [{"type": "audio", "audio": A["c"]}]},
|
||
{"role": "assistant", "content": "<|mixed_start|>"}], True),
|
||
"chat_nothink_c": ([{"role": "user", "content": [{"type": "audio", "audio": A["c"]}]},
|
||
{"role": "assistant", "content": "<think></think>"}], True),
|
||
}
|
||
inputs = {}
|
||
for name, (conv, cont) in convs.items():
|
||
kw = dict(add_generation_prompt=not cont, continue_final_message=cont)
|
||
text = proc.apply_chat_template(conv, tokenize=False, **kw)
|
||
inp = proc.apply_chat_template(conv, tokenize=True, return_dict=True, return_tensors="pt", **kw).to(DEV)
|
||
inputs[name] = inp
|
||
ids = inp["input_ids"][0]
|
||
save(f"prompt/{name}_ids", ids, "prompt token ids (audio placeholders expanded)")
|
||
wav = next((c["audio"] for m in conv if isinstance(m["content"], list) for c in m["content"]), None)
|
||
if wav is not None:
|
||
wn = os.path.basename(wav).replace("_16k.wav", "")
|
||
same = np.array_equal(inp["input_features"][0].cpu().numpy(), feats[wn])
|
||
assert same, f"processor features differ from feature extractor for {name}"
|
||
M["prompts"][name] = {"rendered": text, "len": int(ids.numel()), "n_audio_tokens": int((ids == AUDIO_TOKEN).sum()),
|
||
"ids_file": f"prompt/{name}_ids.npy", "add_generation_prompt": not cont,
|
||
"continue_final_message": cont,
|
||
"conversation": [{**m, "content": [{**c, "audio": os.path.relpath(c["audio"], OUT)} if "audio" in c else c
|
||
for c in m["content"]] if isinstance(m["content"], list) else m["content"]}
|
||
for m in conv],
|
||
"ids": ids.tolist() if ids.numel() < 400 else None}
|
||
|
||
print("backbone", flush=True)
|
||
|
||
def backbone(m, inp):
|
||
with capture({"final_norm": m.model.language_model.embedding_norm}) as g:
|
||
out = m.model(**{k: v for k, v in inp.items() if k != "attention_mask"}, use_cache=False)
|
||
return g["final_norm"], m.lm_head(out.last_hidden_state[0, -1]).float()
|
||
|
||
with torch.no_grad():
|
||
inp = inputs["asr_en_c"]
|
||
hidden, logits = backbone(model, inp)
|
||
src = {"prompt": "prompt/asr_en_c_ids", "audio_features": "enc/c_projector"}
|
||
save("lm/asr_en_c_final_norm", hidden[0], "backbone output after embedding_norm [L, 2048]", **src)
|
||
save("lm/asr_en_c_last_logits", logits, "last-position logits [128000] (tied lm_head)", **src)
|
||
hidden_b, logits_b = backbone(model_bf, inp)
|
||
floor("backbone.final_norm", hidden_b, hidden)
|
||
floor("backbone.last_logits", logits_b, logits, argmax_fp32=int(logits.argmax()), argmax_bf16=int(logits_b.argmax()))
|
||
del hidden, hidden_b
|
||
|
||
for name in ("asr_en_c", "asr_plain_c"):
|
||
inp = inputs[name]
|
||
gen = model.generate(**inp, max_new_tokens=256, **GREEDY).sequences[0, inp["input_ids"].shape[1]:]
|
||
text = proc.decode(gen, skip_special_tokens=True)
|
||
M["results"][f"{name}_greedy_fp32"] = {"text": text, "n_tokens": int(gen.numel()), "ids": gen.tolist()}
|
||
print(f" {name}: {text!r}", flush=True)
|
||
save("gen/asr_en_c_greedy_ids", M["results"]["asr_en_c_greedy_fp32"]["ids"], "greedy ASR generated ids (fp32), ends with <|im_end|>",
|
||
prompt="prompt/asr_en_c_ids")
|
||
|
||
# long sequence: speech.wav (~5.3k positions)
|
||
inp = inputs["asr_plain_speech"]
|
||
L = inp["input_ids"].shape[1]
|
||
pos = sorted({0, 100, 255, 256, 511, 1000, 2047, 2048, 3000, 4096, L // 2, L - 2, L - 1})
|
||
out = model.model(**{k: v for k, v in inp.items() if k != "attention_mask"}, use_cache=False)
|
||
hl = out.last_hidden_state[0, pos]
|
||
ll = model.lm_head(out.last_hidden_state[0, -1]).float()
|
||
del out
|
||
src = {"prompt": "prompt/asr_plain_speech_ids", "audio_features": "enc/speech_projector"}
|
||
save("lm/long_speech_positions", np.array(pos), "positions sampled from the long ASR prompt", **src)
|
||
save("lm/long_speech_final_norm_at_positions", hl, "backbone final-norm hidden at long_speech_positions [P, 2048]", **src)
|
||
save("lm/long_speech_last_logits", ll, "last-position logits of the long ASR prompt [128000]", **src)
|
||
out = model_bf.model(**{k: v for k, v in inp.items() if k != "attention_mask"}, use_cache=False)
|
||
floor("backbone_long.final_norm_at_positions", out.last_hidden_state[0, pos], hl)
|
||
# The C++ long prefill compares the last row only; its floor is about twice the 13-position average.
|
||
floor("backbone_long.final_norm_last", out.last_hidden_state[0, -1], hl[-1])
|
||
llb = model_bf.lm_head(out.last_hidden_state[0, -1]).float()
|
||
floor("backbone_long.last_logits", llb, ll, argmax_fp32=int(ll.argmax()), argmax_bf16=int(llb.argmax()))
|
||
del out
|
||
torch.cuda.empty_cache()
|
||
|
||
def run_gen(name, order_seed=ORDER_SEED, **kw):
|
||
inp = inputs[name]
|
||
L = inp["input_ids"].shape[1]
|
||
with frame_recorder(model, order_seed) as fr, text_logit_recorder(model) as tl, torch.no_grad():
|
||
o = model.generate(**inp, return_audio_eos=True, **GREEDY, **kw)
|
||
kinds, toks, frames, pattern = slots_of(o, L)
|
||
return o, fr, tl, kinds, toks, frames, pattern
|
||
|
||
def dump_gen(tag, name, o, fr, tl, kinds, toks, frames, pattern, settings):
|
||
src = {"prompt": f"prompt/{name}_ids", **settings}
|
||
F = len(fr["codes"])
|
||
save(f"gen/{tag}_slot_kind", kinds, "per generated slot: 0 = text token, 1 = audio frame", **src)
|
||
save(f"gen/{tag}_slot_token", toks, "per slot token id (125032 for audio frames)", **src)
|
||
save(f"gen/{tag}_slot_codes", frames, "per slot codes [S, 12] (-1 for text slots; EOS frame = all 2016)", **src)
|
||
if F:
|
||
save(f"gen/{tag}_frame_hidden", torch.stack(fr["hidden"]), "backbone final-norm hidden fed to the depth decoder per frame [F, 2048]", **src)
|
||
save(f"gen/{tag}_frame_order", torch.stack(fr["order"]), "MaskGIT order per frame [F, 12]; passes reveal order[0:4], [4:8], [8:12]", **src)
|
||
save(f"gen/{tag}_frame_codes_raw", torch.stack(fr["codes"]), "sampled codes per frame before EOS normalization [F, 12]", **src)
|
||
save(f"gen/{tag}_frame_pass_codes", torch.stack(fr["pass_codes"]), "codes input to each depth pass [F, 3, 12] (masked entries irrelevant)", **src)
|
||
save(f"gen/{tag}_frame_pass_masked", torch.stack(fr["pass_masked"]), "masked flags of each depth pass [F, 3, 12] (1 = masked)", **src)
|
||
nl = min(F, MAX_LOGIT_FRAMES)
|
||
save(f"gen/{tag}_frame_pass_logits", torch.stack(fr["pass_logits"][:nl]), f"depth logits of each pass, first {nl} frames [F', 3, 12, 2017]", **src)
|
||
if tl:
|
||
save(f"gen/{tag}_text_top20_vals", torch.stack([t.values for t in tl]), "per text step top-20 logits", **src)
|
||
n_audio = int(o.audio_lengths[0])
|
||
text = proc.decode([t for t in toks.tolist() if t != AUDIO_OUT_TOKEN], skip_special_tokens=True)
|
||
M["results"][tag] = {"pattern": pattern, "n_slots": int(len(kinds)), "n_frames_incl_eos": n_audio, "text": text,
|
||
"settings": settings}
|
||
print(f" {tag}: {pattern} text={text!r}", flush=True)
|
||
|
||
print("generation", flush=True)
|
||
greedy = {**GREEDY, "text": "greedy", "orders": f"injected numpy default_rng({ORDER_SEED}).permutation(12) per frame"}
|
||
o, fr, tl, kinds, toks, frames, pat = run_gen("tts_fox", max_new_tokens=400)
|
||
dump_gen("tts_fox_greedy", "tts_fox", o, fr, tl, kinds, toks, frames, pat, {**greedy, "max_new_tokens": 400})
|
||
tts_frames = frames[kinds == 1]
|
||
tts_codes = torch.as_tensor(tts_frames[~(tts_frames == EOS_CODE).any(1)]).T.contiguous() # [12, T] without EOS
|
||
tts_fr = fr
|
||
|
||
# depth decoder logits noise floor: replay the fp32 pass inputs through the bf16 depth decoder
|
||
with torch.no_grad():
|
||
nl = min(10, len(tts_fr["codes"]))
|
||
lb, lf = [], []
|
||
for i in range(nl):
|
||
h = tts_fr["hidden"][i].to(DEV)
|
||
for p in range(3):
|
||
c, mk = tts_fr["pass_codes"][i][p].to(DEV), tts_fr["pass_masked"][i][p].to(DEV)
|
||
lb.append(model_bf.model.depth_decoder(h[None].bfloat16(), c[None], mk[None])[0].float())
|
||
lf.append(tts_fr["pass_logits"][i][p].to(DEV))
|
||
lb, lf = torch.stack(lb), torch.stack(lf)
|
||
floor("depth.logits", lb, lf, argmax_agree=float((lb.argmax(-1) == lf.argmax(-1)).float().mean()))
|
||
|
||
o, fr, tl, kinds, toks, frames, pat = run_gen("s2s_c", order_seed=S2S_ORDER_SEED, max_new_tokens=1024)
|
||
dump_gen("s2s_c_greedy", "s2s_c", o, fr, tl, kinds, toks, frames, pat,
|
||
{**greedy, "orders": f"injected numpy default_rng({S2S_ORDER_SEED}).permutation(12) per frame", "max_new_tokens": 1024})
|
||
# same prompt, default order seed, cut at max_new_tokens in the middle of the reply (it runs to about 160 slots)
|
||
o, fr, tl, kinds, toks, frames, pat = run_gen("s2s_c", max_new_tokens=96)
|
||
dump_gen("s2s_c_greedy_trunc", "s2s_c", o, fr, tl, kinds, toks, frames, pat, {**greedy, "max_new_tokens": 96})
|
||
|
||
print("codec", flush=True)
|
||
gen_ = torch.Generator(device="cpu").manual_seed(0)
|
||
code_sets = {"tts": tts_codes, "rand200": torch.randint(0, 2016, (12, 200), generator=gen_)}
|
||
for cname, codes in code_sets.items():
|
||
codes = codes.to(DEV)
|
||
save(f"codec/{cname}_codes", codes, "codes [12, T] in [0, 2016) " +
|
||
("(greedy TTS frames, EOS dropped)" if cname == "tts" else "(torch.randint, CPU Generator seed 0)"))
|
||
src = {"codes": f"codec/{cname}_codes"}
|
||
|
||
def codec_mods(c):
|
||
return {"project_out": c.quantizer.project_out, "from_quantizer": c.from_quantizer,
|
||
"dec_final_norm": c.decoder.lfm.embedding_norm, "lin": c.decoder.lin}
|
||
with torch.no_grad():
|
||
with capture(codec_mods(codec)) as g:
|
||
wav = codec.decode(codes[None]).audio_values[0]
|
||
save(f"codec/{cname}_project_out", g["project_out"][0], "quantizer.project_out output [T, 256]", **src)
|
||
save(f"codec/{cname}_from_quantizer", g["from_quantizer"][0], "from_quantizer output [T, 512]", **src)
|
||
save(f"codec/{cname}_dec_final_norm", g["dec_final_norm"][0], "decoder LFM3 output after embedding_norm [6T, 512]", **src)
|
||
nl = min(48, g["lin"].shape[1])
|
||
save(f"codec/{cname}_lin_first", g["lin"][0, :nl], f"decoder.lin output, first {nl} positions [{nl}, 1282] "
|
||
"(log_mag = [:641], angle = [641:])", **src)
|
||
save(f"codec/{cname}_wave", wav, "offline decode, 24 kHz float [1920*T]", **src)
|
||
with capture(codec_mods(codec_bf)) as gb:
|
||
wav_b = codec_bf.decode(codes[None]).audio_values[0]
|
||
if cname == "tts":
|
||
for k in g:
|
||
floor(f"codec.{k}", gb[k], g[k])
|
||
floor(f"codec.{cname}_wave", wav_b, wav)
|
||
T, s, i, parts, sizes = codes.shape[1], 0, 0, [], []
|
||
cache = SCache(codec.config)
|
||
while s < T:
|
||
e_ = min(T, s + (7, 13, 20)[i % 3])
|
||
parts.append(codec.decode(codes[None, :, s:e_], streaming_cache=cache, finalize=e_ == T).audio_values[0])
|
||
sizes.append((e_ - s, parts[-1].numel()))
|
||
s, i = e_, i + 1
|
||
sw = torch.cat(parts)
|
||
save(f"codec/{cname}_stream_chunks", np.array(sizes), "streaming chunks [n, 2] = (frames, output samples); "
|
||
"last chunk finalize=True", **src)
|
||
save(f"codec/{cname}_stream_wave", sw, "concatenated streaming output", **src)
|
||
|
||
print("multiturn", flush=True)
|
||
multiturn(model, proc)
|
||
print("probes", flush=True)
|
||
tokenizer_cases(proc)
|
||
ignore_merges_cases()
|
||
M["meta"]["elapsed_s"] = round(time.time() - t0, 1)
|
||
with open(f"{OUT}/manifest.json", "w") as f:
|
||
json.dump(M, f, indent=1)
|
||
total = sum(os.path.getsize(os.path.join(dp, fn)) for dp, _, fns in os.walk(OUT) for fn in fns)
|
||
print(f"done in {M['meta']['elapsed_s']} s, {len(M['arrays'])} arrays, {total / 1e6:.1f} MB", flush=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|