File size: 1,517 Bytes
f4e441a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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/wav2vec2-large-mms-1b-nhi-adapterft-ilv_fold1"
target_lang = "nhi"
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
|