import torch
import soundfile as sf
from transformers import AutoProcessor, AutoModelForCausalLM
from whisper_normalizer.basic import BasicTextNormalizer

CKPT = "/home/bot/inference"
AUDIO = "/home/bot/inference/audio.wav"

device = "cuda" if torch.cuda.is_available() else "cpu"

# Cargar processor desde tu carpeta
processor = AutoProcessor.from_pretrained(
    CKPT,
    trust_remote_code=True,
    local_files_only=True
)


# Cargar modelo desde tu carpeta
model = AutoModelForCausalLM.from_pretrained(
    CKPT,
    trust_remote_code=True,
    attn_implementation="eager",
    local_files_only=True
).to(device)


model.eval()

audio, sr = sf.read(AUDIO)

user_msg = {"role": "user", "content": "<|audio_1|>\nTranscribe the audio clip into Galician text."}

prompt = processor.tokenizer.apply_chat_template(
    [user_msg],
    tokenize=False,
    add_generation_prompt=True
)

inputs = processor(
    text=prompt,
    audios=[(audio, sr)],
    return_tensors="pt"
)

inputs = {k: v.to(device) for k, v in inputs.items() if v is not None}

L = inputs["input_ids"].shape[1]

gen = model.generate(
    **inputs,
    max_new_tokens=64,
    eos_token_id=processor.tokenizer.eos_token_id,
    num_logits_to_keep=1
)

hyp = processor.batch_decode(gen[:, L:], skip_special_tokens=True)[0]

norm = BasicTextNormalizer()
print(norm(hyp))