| import torch |
| import torch.nn.functional as F |
| import torchaudio |
| from transformers import AutoConfig, Wav2Vec2Processor |
|
|
| from Wav2Vec2ForSpeechClassification import Wav2Vec2ForSpeechClassification |
|
|
| MY_MODEL = "myrun3" |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| config = AutoConfig.from_pretrained(MY_MODEL) |
| processor = Wav2Vec2Processor.from_pretrained(MY_MODEL) |
| sampling_rate = processor.feature_extractor.sampling_rate |
| model = Wav2Vec2ForSpeechClassification.from_pretrained(MY_MODEL).to(device) |
|
|
| def speech_file_to_array_fn(path, sampling_rate): |
| speech_array, _sampling_rate = torchaudio.load(path) |
| resampler = torchaudio.transforms.Resample(_sampling_rate) |
| speech = resampler(speech_array).squeeze().numpy() |
| return speech |
|
|
|
|
| def predict(path, sampling_rate): |
| speech = speech_file_to_array_fn(path, sampling_rate) |
| features = processor(speech, sampling_rate=sampling_rate, return_tensors="pt", padding=True) |
|
|
| input_values = features.input_values.to(device) |
| attention_mask = features.attention_mask.to(device) |
|
|
| with torch.no_grad(): |
| logits = model(input_values, attention_mask=attention_mask).logits |
|
|
| scores = F.softmax(logits, dim=1).detach().cpu().numpy()[0] |
| outputs = [{"Emotion": config.id2label[i], "Score": f"{round(score * 100, 3):.1f}%"} for i, score in enumerate(scores)] |
| return outputs |
|
|
| res = predict("test.wav", 16000) |
| max = max(res, key=lambda x: x['Score']) |
| print("Expected anger:", max) |
|
|
