anyantudre's picture
Update goai_stt.py
98f2613 verified
raw
history blame
1.23 kB
import librosa
import torch
import time
from transformers import set_seed, Wav2Vec2ForCTC, AutoProcessor
def goai_stt(fichier):
"""
Transcrire un fichier audio donné.
Paramètres
----------
fichier: str
Le chemin d'accès au fichier audio.
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