examples/turn_decision.py

"""Evaluate an audio file using the packaged turn-taking question."""

import argparse
import json
from pathlib import Path

import torch
from transformers import AutoModel, AutoProcessor

parser = argparse.ArgumentParser()
parser.add_argument("audio")
parser.add_argument("--model", default=str(Path(__file__).resolve().parents[1]))
args = parser.parse_args()
policy = json.loads((Path(__file__).with_name("turn_policy.json")).read_text())
processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
model = AutoModel.from_pretrained(
    args.model,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="sdpa",
)
inputs = processor(audio=args.audio, state={}, questions=policy["questions"])
inputs = inputs.to(model.device, dtype=torch.bfloat16)
with torch.inference_mode():
    outputs = model(**inputs)
result = processor.post_process(outputs, questions=policy["questions"])
print(json.dumps(result, indent=2))