Update README.md

README.md

@@ -9,7 +9,77 @@ tags:
 pipeline_tag: audio-classification
 ---
 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from transformers import AutoConfig, Wav2Vec2Processor
+import librosa
+import IPython.display as ipd
+import numpy as np
+import pandas as pd
+
 model_name_or_path = "quaja/hubert-base-amharic-speech-emotion-recognition"
 config = AutoConfig.from_pretrained(model_name_or_path)
 feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name_or_path)
-model = HubertForSpeechClassification.from_pretrained(model_name_or_path)
+sampling_rate = feature_extractor.sampling_rate
+model = HubertForSpeechClassification.from_pretrained(model_name_or_path)
+
+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 = feature_extractor(speech, sampling_rate=sampling_rate, return_tensors="pt", padding=True)
+
+    input_values = features.input_values.to(device)
+
+    with torch.no_grad():
+        logits = model(input_values).logits
+
+    scores = F.softmax(logits, dim=1).detach().cpu().numpy()[0]
+    outputs = [{"Label": config.id2label[i], "Score": f"{round(score * 100, 3):.1f}%"} for i, score in enumerate(scores)]
+    return outputs
+
+
+STYLES = """
+<style>
+div.display_data {
+    margin: 0 auto;
+    max-width: 500px;
+}
+table.xxx {
+    margin: 50px !important;
+    float: right !important;
+    clear: both !important;
+}
+table.xxx td {
+    min-width: 300px !important;
+    text-align: center !important;
+}
+</style>
+""".strip()
+
+def prediction(df_row):
+    path, label = df_row["path"], df_row["emotion"]
+    df = pd.DataFrame([{"Emotion": label, "Sentence": "    "}])
+    setup = {
+        'border': 2,
+        'show_dimensions': True,
+        'justify': 'center',
+        'classes': 'xxx',
+        'escape': False,
+    }
+    ipd.display(ipd.HTML(STYLES + df.to_html(**setup) + "<br />"))
+    speech, sr = torchaudio.load(path)
+    resampler = torchaudio.transforms.Resample(sr)
+    speech = resampler(speech[0]).squeeze().numpy()
+    ipd.display(ipd.Audio(data=np.asarray(speech), autoplay=True, rate=sampling_rate))
+
+    outputs = predict(path, sampling_rate)
+    r = pd.DataFrame(outputs)
+    ipd.display(ipd.HTML(STYLES + r.to_html(**setup) + "<br />"))