Views
No views yet
1import torch
2import librosa
3import datasets
4from transformers import Wav2Vec2ForMultiTask, Wav2Vec2Processor
5
6def map_to_array(batch):
7 speech, _ = librosa.load(batch["file"], sr=16000, mono=True)
8 batch["speech"] = speech
9 return batch
10
11def map_to_pred_multitask(batch):
12 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
13 input_values = processor(batch["speech"], sampling_rate=16000, return_tensors="pt", padding="longest").input_values
14 with torch.no_grad():
15 logits = model(input_values.to(device)).logits
16 predicted_ids_ctc = torch.argmax(logits[1], dim=-1)
17 transcription = processor.batch_decode(predicted_ids_ctc)
18 batch["transcription"] = transcription
19 predicted_ids = torch.argmax(logits[0], dim=-1)
20 batch['predictions'] = predicted_ids
21 return batch
22
23processor = Wav2Vec2Processor.from_pretrained(MODEL_PATH)
24model = Wav2Vec2ForMultiTask.from_pretrained(MODEL_PATH)
25
26test_dataset = test_dataset.map(map_to_array)
27result = test_dataset.map(map_to_pred_multitask)1@inproceedings{getman23_slate,
2 author={Yaroslav Getman and Ragheb Al-Ghezi and Tamas Grosz and Mikko Kurimo},
3 title={{Multi-task wav2vec2 Serving as a Pronunciation Training System for Children}},
4 year=2023,
5 booktitle={Proc. 9th Workshop on Speech and Language Technology in Education (SLaTE)},
6 pages={36--40},
7 doi={10.21437/SLaTE.2023-8}
8}