| import librosa | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers import Pipeline, AutoConfig, Wav2Vec2Processor | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model_name_or_path = "KELONMYOSA/wav2vec2-xls-r-300m-emotion-ru" | |
| config = AutoConfig.from_pretrained(model_name_or_path) | |
| processor = Wav2Vec2Processor.from_pretrained(model_name_or_path) | |
| sampling_rate = processor.feature_extractor.sampling_rate | |
| class SpeechEmotionRecognitionPipeline(Pipeline): | |
| def _sanitize_parameters(self, **pipeline_parameters): | |
| return {}, {}, {} | |
| def preprocess(self, audio, second_text=None): | |
| speech, sr = librosa.load(audio, sr=sampling_rate) | |
| features = processor(speech, sampling_rate=sampling_rate, return_tensors="pt", padding=True) | |
| return features.input_values.to(device) | |
| def _forward(self, model_inputs): | |
| return self.model(model_inputs) | |
| def postprocess(self, model_outputs): | |
| logits = model_outputs.logits | |
| scores = F.softmax(logits, dim=1).detach().cpu().numpy()[0] | |
| outputs = [{"label": config.id2label[i], "score": round(score, 5)} for i, score in | |
| enumerate(scores)] | |
| return outputs | |