Update README.md

README.md

@@ -32,7 +32,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchaudio
-from transformers import AutoConfig, Wav2Vec2Processor
+from transformers import AutoConfig, Wav2Vec2FeatureExtractor
 
 import librosa
 import IPython.display as ipd
@@ -44,8 +44,8 @@ import pandas as pd
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_name_or_path = "m3hrdadfi/wav2vec2-xlsr-greek-speech-emotion-recognition"
 config = AutoConfig.from_pretrained(model_name_or_path)
-processor = Wav2Vec2Processor.from_pretrained(model_name_or_path)
-sampling_rate = processor.feature_extractor.sampling_rate
+feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name_or_path)
+sampling_rate = feature_extractor.sampling_rate
 model = Wav2Vec2ForSpeechClassification.from_pretrained(model_name_or_path).to(device)
 ```
 
@@ -59,13 +59,11 @@ def speech_file_to_array_fn(path, sampling_rate):
 
 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)
+    inputs = feature_extractor(speech, sampling_rate=sampling_rate, return_tensors="pt", padding=True)
+    inputs = {key: inputs[key].to(device) for key in inputs}
 
     with torch.no_grad():
-        logits = model(input_values, attention_mask=attention_mask).logits
+        logits = model(**inputs).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)]
@@ -73,17 +71,17 @@ def predict(path, sampling_rate):
 ```
 
 ```python
-path = "/path/to/audio.wav"
+path = "/path/to/disgust.wav"
 outputs = predict(path, sampling_rate)
 ```
 
 ```bash
 [
-	{'Emotion': 'anger', 'Score': '0.0%'},
-	{'Emotion': 'disgust', 'Score': '99.2%'},
-	{'Emotion': 'fear', 'Score': '0.1%'},
-	{'Emotion': 'happiness', 'Score': '0.3%'},
-	{'Emotion': 'sadness', 'Score': '0.5%'}
+\t{'Emotion': 'anger', 'Score': '0.0%'},
+\t{'Emotion': 'disgust', 'Score': '99.2%'},
+\t{'Emotion': 'fear', 'Score': '0.1%'},
+\t{'Emotion': 'happiness', 'Score': '0.3%'},
+\t{'Emotion': 'sadness', 'Score': '0.5%'}
 ]
 ```