"""Load Nawah-ASR-118M-v5. Use load_v5() -- a bare from_pretrained runs at the wrong frame rate.

The audio tower's SHAPE reconstructs correctly from config.json (d_model 576, 12 layers). What
does not survive is the frame rate: Qwen2Audio's encoder ends in a stride-2 average pool that
v5 removes, and that pooler has no parameters and is rebuilt on every load. So from_pretrained
returns all 300 tensors byte-identical, no warning, at 25 Hz instead of the 50 Hz it was
trained at. config.json records audio_config.pool_stride: 1; stock transformers ignores it.

    from load_model import load_v5, transcribe
    model, tok, fe = load_v5("oddadmix/Nawah-ASR-118M-v5")
    print(transcribe(model, tok, fe, wav))      # float32 mono @ 16 kHz
"""
import torch
from transformers import (AutoTokenizer, Qwen2AudioForConditionalGeneration,
                          WhisperFeatureExtractor)
from build_model import set_audio_frame_rate


def load_v5(repo: str, token: str | None = None, dtype=torch.float32):
    model = Qwen2AudioForConditionalGeneration.from_pretrained(repo, dtype=dtype, token=token)
    set_audio_frame_rate(model, 1)                     # 1 = 50 Hz
    model.eval()
    _, n = model.model.audio_tower._get_feat_extract_output_lengths(torch.tensor([3000]))
    assert int(n) == 1500, f"expected 50 Hz (1500 tokens / 30 s), got {int(n)}"
    tok = AutoTokenizer.from_pretrained(repo, token=token)
    fe = WhisperFeatureExtractor.from_pretrained(repo, token=token)
    return model, tok, fe


@torch.no_grad()
def transcribe(model, tok, fe, wav, sr=16000, max_new_tokens=160):
    feats = fe([wav], sampling_rate=sr, return_attention_mask=True, return_tensors="pt")
    _, n = model.model.audio_tower._get_feat_extract_output_lengths(feats.attention_mask.sum(-1))
    cv = tok.convert_tokens_to_ids
    # The placeholder is "<audio>" (id 32000). "<|AUDIO|>" is NOT in this tokenizer -- it maps
    # to UNK (id 3), which silently fills the prefix with unknown tokens and transcribes noise.
    aud, a0, a1 = cv("<audio>"), cv("<|audio_start|>"), cv("<|audio_end|>")
    assert tok.unk_token_id not in (aud, a0, a1), "audio special tokens missing from tokenizer"
    prefix = [tok.bos_token_id, a0] + [aud] * int(n[0]) + [a1]
    ids = torch.tensor([prefix])
    out = model.generate(input_ids=ids, attention_mask=torch.ones_like(ids),
                         input_features=feats.input_features,
                         feature_attention_mask=feats.attention_mask,
                         max_new_tokens=max_new_tokens, do_sample=False)
    return tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True).strip()