File size: 1,272 Bytes
ffa317c |
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 43 |
"""
Inference main class.
Author: Marcely Zanon Boito, 2024
"""
from .CTC_model import mHubertForCTC
import torch
from transformers import Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor, Wav2Vec2Processor
from transformers import HubertConfig
from datasets import load_dataset
fbk_test_id = 'FBK-MT/Speech-MASSIVE-test'
mhubert_id = 'utter-project/mHuBERT-147'
def load_asr_model():
# Load the ASR model
tokenizer = Wav2Vec2CTCTokenizer("asr/vocab.json", unk_token="[UNK]", pad_token="[PAD]", word_delimiter_token="|")
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(mhubert_id)
processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
config = HubertConfig.from_pretrained("naver/mHuBERT-147-ASR-fr")
model = mHubertForCTC.from_pretrained("naver/mHuBERT-147-ASR-fr", config=config)
model.eval()
return model, processor
def run_asr_inference(model, processor, example):
audio = processor(example["array"], sampling_rate=example["sampling_rate"]).input_values[0]
input_values = torch.tensor(audio).unsqueeze(0)
with torch.no_grad():
logits = model(input_values).logits
pred_ids = torch.argmax(logits, dim=-1)
prediction = processor.batch_decode(pred_ids)[0].replace('[CTC]', "")
return prediction
|