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
262 lines
9.4 KiB
Python
262 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate pinned Fun-ASR-Nano prompt and Qwen decoder references."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import subprocess
|
|
import sys
|
|
from array import array
|
|
from pathlib import Path
|
|
|
|
|
|
TRANSFORMERS_COMMIT = "f9966442ac24fff57060774ce22e1884760f4a3b"
|
|
TRANSFORMERS_SOURCE_SHA256 = {
|
|
"src/transformers/models/fun_asr_nano/configuration_fun_asr_nano.py": (
|
|
"174abb3aad5b38f7208c7eae140531be900eafb3a5d7f9f5562c300c93240e2d"
|
|
),
|
|
"src/transformers/models/fun_asr_nano/modeling_fun_asr_nano.py": (
|
|
"70f8230fd5a75e8a119cea1c8dc919c3f51d8aa7c5d8f1c23c5402b856ebebfb"
|
|
),
|
|
"src/transformers/models/fun_asr_nano/processing_fun_asr_nano.py": (
|
|
"bc1601a59f6acef6e5be743e511891fad688afa192d10e722fdea2dc42b5f72a"
|
|
),
|
|
"src/transformers/models/qwen3/modeling_qwen3.py": (
|
|
"966626c5f2c0f729dd6b61b99c000038602805f7d7ad73957507a63f8c2de563"
|
|
),
|
|
}
|
|
MODEL_REVISION = "854d88f94205cd17d2afdb24332130d86fbe654a"
|
|
MODEL_CONFIG_SHA256 = "c7c4a30316929631ac5fabc5fb3c0dd3278dcc9809670720c5920186285d004a"
|
|
MODEL_SAFETENSORS_SHA256 = (
|
|
"335ca3e74917f1156690400e2c344350112950165789cf78ce3d0a367affd821"
|
|
)
|
|
ITN_PROMPT = "语音转写："
|
|
NO_ITN_PROMPT = "语音转写，不进行文本规整："
|
|
AUDIO_TOKENS = 2
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--transformers-src", type=Path, required=True)
|
|
parser.add_argument("--model-dir", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--device", default="cuda")
|
|
return parser.parse_args()
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def require_revision(transformers_src: Path) -> None:
|
|
actual = subprocess.check_output(
|
|
["git", "rev-parse", "HEAD"], cwd=transformers_src, text=True
|
|
).strip()
|
|
if actual != TRANSFORMERS_COMMIT:
|
|
raise RuntimeError(
|
|
f"expected Transformers {TRANSFORMERS_COMMIT}, found {actual}"
|
|
)
|
|
dirty = subprocess.check_output(
|
|
["git", "status", "--porcelain", "--untracked-files=no"],
|
|
cwd=transformers_src,
|
|
text=True,
|
|
).strip()
|
|
if dirty:
|
|
raise RuntimeError("Transformers reference checkout has tracked modifications")
|
|
for relative_path, expected_hash in TRANSFORMERS_SOURCE_SHA256.items():
|
|
actual_hash = sha256(transformers_src / relative_path)
|
|
if actual_hash != expected_hash:
|
|
raise RuntimeError(
|
|
f"unexpected Transformers source hash for {relative_path}: "
|
|
f"{actual_hash}"
|
|
)
|
|
|
|
|
|
def chat_text(prompt: str) -> str:
|
|
return (
|
|
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
|
f"<|im_start|>user\n{prompt}<|object_ref_start|><|im_end|>\n"
|
|
"<|im_start|>assistant\n"
|
|
)
|
|
|
|
|
|
def expanded_prompt(
|
|
tokenizer, prompt: str, audio_token_id: int
|
|
) -> tuple[list[int], list[int]]:
|
|
ids = tokenizer.encode(chat_text(prompt), add_special_tokens=False)
|
|
if ids.count(audio_token_id) != 1:
|
|
raise RuntimeError("expected one audio placeholder in Fun-ASR-Nano prompt")
|
|
expanded: list[int] = []
|
|
positions: list[int] = []
|
|
for token_id in ids:
|
|
if token_id == audio_token_id:
|
|
for _ in range(AUDIO_TOKENS):
|
|
positions.append(len(expanded))
|
|
expanded.append(token_id)
|
|
else:
|
|
expanded.append(token_id)
|
|
return expanded, positions
|
|
|
|
|
|
def deterministic_audio(torch):
|
|
values = [
|
|
((index % 97) - 48) / 97.0 + math.sin(index * 0.013) * 0.1
|
|
for index in range(AUDIO_TOKENS * 1024)
|
|
]
|
|
return torch.tensor(values, dtype=torch.float32).reshape(AUDIO_TOKENS, 1024)
|
|
|
|
|
|
def append_f32(blob: array, values) -> dict[str, int]:
|
|
flat = values.detach().to(device="cpu", dtype=values.new_zeros(()).float().dtype)
|
|
flat = flat.contiguous().view(-1).tolist()
|
|
offset = len(blob)
|
|
blob.extend(flat)
|
|
return {"offset_f32": offset, "count": len(flat)}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
require_revision(args.transformers_src)
|
|
sys.path.insert(0, str(args.transformers_src / "src"))
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
from transformers import AutoProcessor, FunAsrNanoForConditionalGeneration
|
|
|
|
torch.manual_seed(0)
|
|
torch.set_grad_enabled(False)
|
|
torch.set_float32_matmul_precision("highest")
|
|
|
|
config_path = args.model_dir / "config.json"
|
|
model_path = args.model_dir / "model.safetensors"
|
|
if sha256(config_path) != MODEL_CONFIG_SHA256:
|
|
raise RuntimeError("unexpected Fun-ASR-Nano config hash")
|
|
if sha256(model_path) != MODEL_SAFETENSORS_SHA256:
|
|
raise RuntimeError("unexpected Fun-ASR-Nano model hash")
|
|
|
|
processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True)
|
|
model = FunAsrNanoForConditionalGeneration.from_pretrained(
|
|
args.model_dir,
|
|
local_files_only=True,
|
|
torch_dtype=torch.float32,
|
|
attn_implementation="eager",
|
|
).to(args.device)
|
|
model.eval()
|
|
|
|
audio_token_id = model.config.audio_token_id
|
|
itn_ids, itn_positions = expanded_prompt(
|
|
processor.tokenizer, ITN_PROMPT, audio_token_id
|
|
)
|
|
no_itn_ids, no_itn_positions = expanded_prompt(
|
|
processor.tokenizer, NO_ITN_PROMPT, audio_token_id
|
|
)
|
|
if len(itn_positions) != AUDIO_TOKENS or len(no_itn_positions) != AUDIO_TOKENS:
|
|
raise RuntimeError("expanded audio token count mismatch")
|
|
|
|
input_ids = torch.tensor([itn_ids], dtype=torch.long, device=args.device)
|
|
attention_mask = torch.ones_like(input_ids)
|
|
audio = deterministic_audio(torch).to(args.device)
|
|
inputs_embeds = model.get_input_embeddings()(input_ids)
|
|
inputs_embeds[0, itn_positions] = audio
|
|
|
|
outputs = model(
|
|
inputs_embeds=inputs_embeds,
|
|
attention_mask=attention_mask,
|
|
use_cache=True,
|
|
logits_to_keep=1,
|
|
return_dict=True,
|
|
)
|
|
logits = [outputs.logits[0, -1].float()]
|
|
greedy_tokens = [int(torch.argmax(logits[0]).item())]
|
|
past_key_values = outputs.past_key_values
|
|
|
|
for _ in range(2):
|
|
step_id = torch.tensor(
|
|
[[greedy_tokens[-1]]], dtype=torch.long, device=args.device
|
|
)
|
|
attention_mask = torch.ones(
|
|
(1, attention_mask.shape[1] + 1), dtype=torch.long, device=args.device
|
|
)
|
|
outputs = model(
|
|
input_ids=step_id,
|
|
attention_mask=attention_mask,
|
|
past_key_values=past_key_values,
|
|
use_cache=True,
|
|
logits_to_keep=1,
|
|
return_dict=True,
|
|
)
|
|
past_key_values = outputs.past_key_values
|
|
logits.append(outputs.logits[0, -1].float())
|
|
greedy_tokens.append(int(torch.argmax(logits[-1]).item()))
|
|
|
|
with safe_open(model_path, framework="pt", device="cpu") as handle:
|
|
text_names = sorted(
|
|
name for name in handle.keys() if name.startswith("model.language_model.")
|
|
)
|
|
catalog_hash = hashlib.sha256("\n".join(text_names).encode()).hexdigest()
|
|
|
|
blob = array("f")
|
|
audio_descriptor = append_f32(blob, audio)
|
|
audio_descriptor["shape"] = [AUDIO_TOKENS, 1024]
|
|
logit_descriptors = []
|
|
for values in logits:
|
|
descriptor = append_f32(blob, values)
|
|
descriptor["shape"] = [151936]
|
|
logit_descriptors.append(descriptor)
|
|
if sys.byteorder != "little":
|
|
blob.byteswap()
|
|
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
data_path = args.output.with_suffix(".bin")
|
|
with data_path.open("wb") as handle:
|
|
blob.tofile(handle)
|
|
|
|
metadata = {
|
|
"schema_version": 1,
|
|
"transformers_commit": TRANSFORMERS_COMMIT,
|
|
"transformers_source_sha256": TRANSFORMERS_SOURCE_SHA256,
|
|
"model_revision": MODEL_REVISION,
|
|
"model_config_sha256": MODEL_CONFIG_SHA256,
|
|
"model_safetensors_sha256": MODEL_SAFETENSORS_SHA256,
|
|
"tokenizer_json_sha256": sha256(args.model_dir / "tokenizer.json"),
|
|
"tokenizer_config_sha256": sha256(args.model_dir / "tokenizer_config.json"),
|
|
"text_weight_catalog_count": len(text_names),
|
|
"text_weight_catalog_sha256": catalog_hash,
|
|
"torch_version": torch.__version__,
|
|
"device": str(args.device),
|
|
"data_file": data_path.name,
|
|
"data_format": "little-endian-float32",
|
|
"audio_token_id": audio_token_id,
|
|
"im_start_token_id": processor.tokenizer.convert_tokens_to_ids("<|im_start|>"),
|
|
"im_end_token_id": processor.tokenizer.convert_tokens_to_ids("<|im_end|>"),
|
|
"eos_token_id": model.config.text_config.eos_token_id,
|
|
"itn": {
|
|
"prompt": ITN_PROMPT,
|
|
"input_ids": itn_ids,
|
|
"audio_token_positions": itn_positions,
|
|
},
|
|
"no_itn": {
|
|
"prompt": NO_ITN_PROMPT,
|
|
"input_ids": no_itn_ids,
|
|
"audio_token_positions": no_itn_positions,
|
|
},
|
|
"audio_embeddings": audio_descriptor,
|
|
"greedy_token_ids": greedy_tokens,
|
|
"generated_text": processor.tokenizer.decode(
|
|
greedy_tokens, skip_special_tokens=True
|
|
),
|
|
"logits": logit_descriptors,
|
|
"data_sha256": sha256(data_path),
|
|
}
|
|
args.output.write_text(json.dumps(metadata, indent=2, ensure_ascii=False) + "\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|