| 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) | |