File size: 2,173 Bytes
763071e |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
"""
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():
def init_config():
config = HubertConfig.from_pretrained(mhubert_id)
config.pad_token_id = processor.tokenizer.pad_token_id
config.ctc_token_id = processor.tokenizer.convert_tokens_to_ids('[CTC]')
config.vocab_size = len(processor.tokenizer)
config.output_hidden_states = False
config.add_interface = True
config.num_interface_layers = 3
return config
# Load the ASR model
tokenizer = Wav2Vec2CTCTokenizer('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 = init_config()
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
if __name__ == '__main__':
# Load the dataset in streaming mode
dataset = load_dataset(fbk_test_id, 'fr-FR', streaming=True)
dataset = dataset['test']
generator = iter(dataset)
# load model
model, processor = load_asr_model()
print(model)
# decode 10 examples from speech-MASSIVE
num_examples= 10
while num_examples >= 0:
example = next(generator)
prediction = run_inference(model, processor, example['audio'])
gold_standard = example['utt']
print("Gold standard:", gold_standard)
print("Prediction:", prediction)
print()
num_examples-=1
|