pipeline

config.json

@@ -10,7 +10,7 @@
   ],
   "attention_dropout": 0.1,
   "auto_map": {
-    "AutoModelForAudioClassification": "emotion_model.Wav2Vec2ForSpeechClassification"
+    "AutoModelForAudioClassification": "KELONMYOSA/wav2vec2-xls-r-300m-emotion-ru--emotion_model.Wav2Vec2ForSpeechClassification"
   },
   "bos_token_id": 1,
   "classifier_proj_size": 256,
@@ -46,6 +46,23 @@
   ],
   "ctc_loss_reduction": "mean",
   "ctc_zero_infinity": true,
+  "custom_pipelines": {
+    "audio-classification": {
+      "default": {
+        "model": {
+          "pt": [
+            "KELONMYOSA/wav2vec2-xls-r-300m-emotion-ru",
+            "main"
+          ]
+        }
+      },
+      "impl": "emotion_recognition_pipeline.SpeechEmotionRecognitionPipeline",
+      "pt": [
+        "Wav2Vec2ForSpeechClassification"
+      ],
+      "tf": []
+    }
+  },
   "diversity_loss_weight": 0.1,
   "do_stable_layer_norm": true,
   "eos_token_id": 2,

emotion_recognition_pipeline.py

@@ -0,0 +1,32 @@
+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