from transformers import Wav2Vec2ForCTC, AutoProcessor
import torchaudio
import torch
import os
import librosa

hf_token = os.getenv("HUGGING_FACE_HUB_TOKEN")

def read_audio_data(file):
    speech_array, sampling_rate = torchaudio.load(file, normalize = True)
    return speech_array, sampling_rate

def load_model():
    model_id = "Lguyogiro/w-apostrophe_wav2vec2-large-mms-1b-oji-adapterft"
    target_lang = "oji"
    processor = AutoProcessor.from_pretrained(model_id, target_lang=target_lang, use_auth_token=hf_token)
    model = Wav2Vec2ForCTC.from_pretrained(model_id, target_lang=target_lang, ignore_mismatched_sizes=True, use_safetensors=True, use_auth_token=hf_token)
    return processor, model


def inference(processor, model, audio_path):
    audio, sampling_rate = librosa.load(audio_path, sr=16000)  # Ensure the correct sampling rate
    inputs = processor(audio, sampling_rate=sampling_rate, return_tensors="pt", padding=True)
    
    with torch.no_grad():
        logits = model(inputs.input_values).logits
    
    # Decode predicted tokens
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]


    #arr, rate = read_audio_data(audio_path)
    #inputs = processor(arr.squeeze().numpy(), sampling_rate=16_000, return_tensors="pt")

    #with torch.no_grad():
    #    outputs = model(**inputs).logits
    #ids = torch.argmax(outputs, dim=-1)[0]
    #transcription = processor.decode(ids)

    return transcription