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
150 lines
7.2 KiB
Python
150 lines
7.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the official Ming3 Audio Thinker and preserve its AR input/output boundary."""
|
|
|
|
import argparse
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
import soundfile as sf
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
from transformers import AutoModelForCausalLM, AutoProcessor
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools/audiocpp_cli"))
|
|
from run_audiocpp_cli_path_tests import ResourceMonitor
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "reference/Ming-3.0-audio-flash/examples"))
|
|
from transformers_audio import build_messages, process_messages
|
|
from types import SimpleNamespace
|
|
|
|
|
|
@torch.inference_mode()
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-dir", type=Path, required=True)
|
|
parser.add_argument("--requests", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--capture", action="store_true",
|
|
help="Save processor tensors, projected audio, AR prompt embeddings and generated IDs")
|
|
parser.add_argument("--attention", choices=("eager", "sdpa"), default="sdpa")
|
|
parser.add_argument("--log", action="store_true", required=True)
|
|
args = parser.parse_args()
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
requests = json.loads(args.requests.read_text(encoding="utf-8"))
|
|
if isinstance(requests, dict):
|
|
requests = requests["requests"]
|
|
processor = AutoProcessor.from_pretrained(args.model_dir, trust_remote_code=True)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
args.model_dir, trust_remote_code=True, dtype=torch.bfloat16,
|
|
device_map="cuda:0", attn_implementation=args.attention).eval()
|
|
manifest = {
|
|
"torch": torch.__version__,
|
|
"gpu": torch.cuda.get_device_name(0),
|
|
"dtype": str(model.dtype),
|
|
"attention": args.attention,
|
|
"config_sha256": hashlib.sha256((args.model_dir / "config.json").read_bytes()).hexdigest(),
|
|
"requests": [],
|
|
"parameter_dtypes": {name: str(model.get_parameter(name).dtype) for name in (
|
|
"model.word_embeddings.weight", "lm_head.weight",
|
|
"model.layers.0.input_layernorm.weight", "model.layers.0.attention.A_log",
|
|
"model.layers.0.attention.dt_bias", "model.layers.1.mlp.gate.weight")},
|
|
"buffer_dtypes": {"model.layers.1.mlp.gate.expert_bias":
|
|
str(model.get_buffer("model.layers.1.mlp.gate.expert_bias").dtype)},
|
|
}
|
|
monitor = ResourceMonitor(os.getpid(), 0, 50)
|
|
stop = threading.Event()
|
|
|
|
def sample_resources():
|
|
while not stop.is_set():
|
|
monitor.sample()
|
|
stop.wait(0.05)
|
|
|
|
executor = ThreadPoolExecutor(max_workers=1)
|
|
sampling_job = executor.submit(sample_resources)
|
|
try:
|
|
for request in requests:
|
|
request = {**request, **request.get("options", {})}
|
|
name = request.get("name", request.get("id"))
|
|
if not name:
|
|
raise ValueError("Each request requires a name or id")
|
|
torch.manual_seed(request.get("seed", 42))
|
|
torch.cuda.synchronize()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
start = time.perf_counter()
|
|
audio = Path(request["audio"]) if "audio" in request else None
|
|
if audio:
|
|
chat = "system_prompt" in request
|
|
messages = build_messages(SimpleNamespace(
|
|
task="chat" if chat else "asr", prompt_as_system=False,
|
|
prompt=request.get("prompt", None if chat else processor.ASR_PROMPT),
|
|
system_prompt=request.get("system_prompt", "")), str(audio.resolve()))
|
|
inputs = process_messages(processor, messages, request.get("enable_thinking", False))
|
|
info = sf.info(audio)
|
|
duration = info.frames / info.samplerate
|
|
else:
|
|
prompt = processor.tokenizer.apply_chat_template(
|
|
[{"role": "user", "content": request["text"]}],
|
|
tokenize=False, add_generation_prompt=True, enable_thinking=False)
|
|
inputs = processor.tokenizer(prompt, return_tensors="pt")
|
|
inputs.pop("token_type_ids", None)
|
|
duration = None
|
|
inputs = inputs.to(model.device)
|
|
if "audio_feats" in inputs:
|
|
inputs["audio_feats"] = inputs["audio_feats"].to(model.dtype)
|
|
tensors = {key: value.detach().cpu().contiguous() for key, value in inputs.items()}
|
|
capture_handle = None
|
|
if args.capture:
|
|
def capture_prefill(module, positional, keywords):
|
|
if "prompt_embeddings" in tensors:
|
|
return
|
|
embeddings = keywords["inputs_embeds"]
|
|
tensors["prompt_embeddings"] = embeddings.detach().cpu().contiguous()
|
|
if audio:
|
|
mask = inputs["input_ids"] == model.config.audio_token_id
|
|
tensors["audio_embeddings"] = embeddings[mask].reshape(
|
|
1, -1, embeddings.shape[-1]).detach().cpu().contiguous()
|
|
|
|
capture_handle = model.model.register_forward_pre_hook(capture_prefill, with_kwargs=True)
|
|
try:
|
|
output = model.generate(**inputs, max_new_tokens=request.get("max_tokens", 256), do_sample=False)
|
|
finally:
|
|
if capture_handle is not None:
|
|
capture_handle.remove()
|
|
torch.cuda.synchronize()
|
|
wall_ms = (time.perf_counter() - start) * 1000
|
|
monitor.sample()
|
|
generated = output[:, inputs["input_ids"].shape[1]:]
|
|
text = processor.decode(generated[0], skip_special_tokens=True).strip()
|
|
tensors["generated_ids"] = generated.cpu().contiguous()
|
|
(args.output_dir / f"{name}.txt").write_text(text + "\n", encoding="utf-8")
|
|
if args.capture:
|
|
save_file(tensors, str(args.output_dir / f"{name}.safetensors"))
|
|
metrics = {
|
|
"name": name, "text": text, "session.wall_ms": wall_ms,
|
|
"input_tokens": inputs["input_ids"].shape[1], "generated_tokens": generated.shape[1],
|
|
"peak_allocated_vram_bytes": torch.cuda.max_memory_allocated(),
|
|
"peak_reserved_vram_bytes": torch.cuda.max_memory_reserved(),
|
|
"audio_duration_sec": duration,
|
|
"rtf": wall_ms / (duration * 1000) if duration else None,
|
|
"instrumented": args.capture,
|
|
"capture_boundary": "actual_thinker_prefill_inputs" if args.capture else None,
|
|
**monitor.metrics.to_json(),
|
|
}
|
|
manifest["requests"].append(metrics)
|
|
print(json.dumps(metrics, ensure_ascii=False), flush=True)
|
|
(args.output_dir / "metrics.json").write_text(
|
|
json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
finally:
|
|
stop.set()
|
|
sampling_job.result()
|
|
executor.shutdown()
|
|
monitor.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|