"""Loader for Nawah-ASR-159M-v1.

DO NOT use Qwen2AudioForConditionalGeneration.from_pretrained on this repo. Two things
break it, and the second one breaks it SILENTLY:

  1. The projector is not the stock single Linear -- it is LayerNorm + 2-layer MLP, so
     the module shapes differ.
  2. This graft nests a full ForCausalLM under model.language_model, so save_pretrained
     emits keys like `language_model.model.model.layers.0...` while from_pretrained
     expects `model.language_model.layers.0...`. The mismatch is reported as MISSING and
     those tensors are then RANDOMLY INITIALISED -- the model loads without error and
     emits fluent nonsense.

So the weights ship as a raw state_dict (weights.pt) and are loaded into a shell built
by build_model.build(), which is exactly how training and evaluation construct it.

    from load_model import load_nawah, transcribe
    model, tok, fe = load_nawah("oddadmix/Nawah-ASR-159M-v1", token=...)
    print(transcribe(model, tok, fe, wav_16k_mono))
"""
from __future__ import annotations

import os

import torch
from transformers import AutoTokenizer, WhisperFeatureExtractor

ENCODER = "oddadmix/whisper-small-arabic-dialectal"


def load_nawah(repo: str, token: str | None = None, dtype=torch.float32, encoder: str = ENCODER):
    from build_model import build
    model, _ = build(encoder, projector="mlp", dtype=dtype)
    if os.path.isdir(repo):
        weights = os.path.join(repo, "weights.pt")
    else:
        from huggingface_hub import hf_hub_download
        weights = hf_hub_download(repo, "weights.pt", token=token)
    model.load_state_dict(torch.load(weights, map_location="cpu"))
    model.eval()
    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
    ids = torch.tensor([[tok.bos_token_id, cv("<|audio_start|>")]
                        + [cv("<audio>")] * int(n[0]) + [cv("<|audio_end|>")]])
    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, num_beams=1,
                         pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)
    return tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True).strip()