Views
No views yet
1import torchaudio
2from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
3import re
4
5def clean_up(transcription):
6 hangul = re.compile('[^ ㄱ-ㅣ가-힣]+')
7 result = hangul.sub('', transcription)
8 return result
9
10model_name "hyyoka/wav2vec2-xlsr-korean-senior"
11processor = Wav2Vec2Processor.from_pretrained(model_name)
12model = Wav2Vec2ForCTC.from_pretrained(model_name)
13speech_array, sampling_rate = torchaudio.load(wav_file)
14feat = processor(speech_array[0],
15 sampling_rate=16000,
16 padding=True,
17 max_length=800000,
18 truncation=True,
19 return_attention_mask=True,
20 return_tensors="pt",
21 pad_token_id=49
22 )
23input = {'input_values': feat['input_values'],'attention_mask':feat['attention_mask']}
24
25outputs = model(**input, output_attentions=True)
26logits = outputs.logits
27predicted_ids = logits.argmax(axis=-1)
28transcription = processor.decode(predicted_ids[0])
29stt_result = clean_up(transcription)