scripts/process_teacher.py

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import torch

from scripts.build_teacher_dataset import Aligner
from speech_bridge_gemma.qwen3_tts_tokenizer_smoke import load_qwen3_codec, qwen3_codes_to_qt


def read_db(db_path: Path) -> list[dict]:
    rows = []
    if db_path.exists():
        for line in db_path.read_text(encoding="utf-8").splitlines():
            if line.strip():
                rows.append(json.loads(line))
    return rows


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dir", required=True)
    parser.add_argument("--target", type=int, default=2500)
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()

    out = Path(args.dir)
    (out / "data").mkdir(parents=True, exist_ok=True)
    codec = load_qwen3_codec("Qwen/Qwen3-TTS-Tokenizer-12Hz", args.device)
    aligner = Aligner(args.device)

    done = 0
    fails = 0
    idle = 0
    while True:
        rows = read_db(out / "db.jsonl")
        pending = [r for r in rows if not (out / "data" / f"{r['id']}.pt").exists()]
        if not pending:
            if done >= args.target or idle >= 360:
                break
            idle += 1
            time.sleep(10)
            continue
        idle = 0
        for r in pending:
            wav = out / "wavs" / f"{r['id']}.wav"
            if not wav.exists():
                continue
            try:
                codes = qwen3_codes_to_qt(codec.encode(str(wav)))
                align = aligner.align(str(wav), r["answer"], int(codes.shape[1]))
            except Exception as exc:
                print(json.dumps({"event": "err", "id": r["id"], "err": str(exc)[:120]}), flush=True)
                fails += 1
                continue
            if align is None:
                fails += 1
                continue
            torch.save({"codes": codes, "align": align.long()}, out / "data" / f"{r['id']}.pt")
            done += 1
        processed = len(list((out / "data").glob("*.pt")))
        print(json.dumps({"event": "progress", "processed": processed, "fails": fails}), flush=True)

    print(json.dumps({"event": "done", "processed": len(list((out / "data").glob("*.pt"))), "fails": fails}), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())