| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
|
|
| from speech_bridge_gemma.qwen3_tts_tokenizer_smoke import decode_qwen3_codes, load_qwen3_codec, save_audio |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--manifest", required=True) |
| parser.add_argument("--out-dir", required=True) |
| parser.add_argument("--codec-model", default="Qwen/Qwen3-TTS-Tokenizer-12Hz") |
| parser.add_argument("--codec-device", default="cuda") |
| return parser.parse_args() |
|
|
|
|
| def load_codes(path: str): |
| codes = torch.load(path, map_location="cpu") |
| if isinstance(codes, dict): |
| codes = codes.get("codes") or codes.get("audio_codes") |
| return torch.as_tensor(codes).detach().cpu().long() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| out_dir = Path(args.out_dir) |
| manifest = json.loads(Path(args.manifest).read_text(encoding="utf-8")) |
| codec = load_qwen3_codec(args.codec_model, args.codec_device) |
| asr_items = [] |
| for item in manifest["items"]: |
| codes_path = item.get("out_codes") |
| if not codes_path: |
| continue |
| codes = load_codes(codes_path) |
| wav, sr = decode_qwen3_codes(codec, codes) |
| wav_path = out_dir / f"{Path(codes_path).stem.replace('_codes', '')}.wav" |
| save_audio(wav_path, wav, sr) |
| item["prediction"] = str(wav_path) |
| item["sample_rate"] = int(sr) |
| asr_items.append({"id": item["id"], "wav": str(wav_path), "expected": item["answer"]}) |
| print(json.dumps({"event": "decoded", "id": item["id"], "wav": str(wav_path), "frames": int(codes.shape[-1])}, ensure_ascii=False), flush=True) |
| (out_dir / "decoded_manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") |
| (out_dir / "asr_items.json").write_text(json.dumps(asr_items, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") |
| print(json.dumps({"event": "done", "decoded": len(asr_items)}, ensure_ascii=False), flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
