File size: 1,424 Bytes
57ea1c7 |
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 |
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
from optimum.bettertransformer import BetterTransformer
import torch
import librosa
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
# os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = '1'
# os.environ['TRANSFORMERS_VERBOSITY'] = 'error'
torch.random.manual_seed(0);
# protobuf==3.20.0
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
model = BetterTransformer.transform(model)
def load_audio(audio_path, processor):
audio, sr = librosa.load(audio_path, sr=None)
input_values = processor(audio, sampling_rate=16000, return_tensors="pt").input_features
return input_values
@torch.inference_mode()
def get_emissions(input_values, model):
results = model(input_values,).logits
return results
def score_audio(audio_path, true_result):
true_result = true_result.split('/')
input_values = load_audio(audio_path, processor)
logits = get_emissions(input_values, model).cpu()
predicted_ids = torch.argmax(logits, dim=-1)
transcription = processor.batch_decode(predicted_ids)[0].lower()
result = {'transcription': transcription,
'score': int(any([x in transcription for x in true_result])),
}
return result
|