import librosa
import torch
import time
from transformers import set_seed, Wav2Vec2ForCTC, AutoProcessor


def goai_stt(fichier, device):
    """
    Transcrire un fichier audio donné.
    
    Paramètres
    ---------- 
    fichier: str
        Le chemin d'accès au fichier audio.

    device: str
        GPU ou CPU
        
    Return
    ---------- 
    transcript: str
        Le texte transcrit.
    """
    
    
    ### assurer reproducibilité
    set_seed(2024) 
    
    start_time = time.time()
    
    ### charger le modèle de transcription 
    model_id = "anyantudre/wav2vec2-large-mms-1b-mos-V1"

    processor = AutoProcessor.from_pretrained(model_id)
    model = Wav2Vec2ForCTC.from_pretrained(model_id, target_lang="mos", device=device, ignore_mismatched_sizes=True)

    ### preprocessing de l'audio
    signal, sampling_rate =  librosa.load(fichier, sr=16000)
    inputs = processor(signal, sampling_rate=16_000, return_tensors="pt", padding=True)
    
    ### faire l'inference
    with torch.no_grad():
        outputs = model(**inputs).logits
    
    pred_ids = torch.argmax(outputs, dim=-1)[0]
    transcription = processor.decode(pred_ids)

    print("Temps écoulé: ", int(time.time() - start_time), " secondes")
    return transcription