Add emotion interface

interfaces/emotion.py

@@ -7,21 +7,20 @@ from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
 from huggingface_hub import HfApi
 
-from label_dicts import MANIFESTO_LABEL_NAMES
+from label_dicts import EMOTION9_LABEL_NAMES
 
 HF_TOKEN = os.environ["hf_read"]
 
 languages = [
-    "Czech", "English", "French", "German", "Hungarian", "Polish", "Slovak"
+    "Czech", "English", "German", "Hungarian", "Polish", "Slovak"
 ]
 domains = {
     "parliamentary speech": "parlspeech",
 }
 
 def build_huggingface_path(language: str):
-    if language == "Czech" or language == "Slovak":
-        return "visegradmedia-emotion/Emotion_RoBERTa_pooled_V4"
-    return "poltextlab/xlm-roberta-large-pooled-MORES"
+    language = language.lower()
+    return f"poltextlab/xlm-roberta-large-pooled-{language}-emotions9"
 
 def predict(text, model_id, tokenizer_id):
     device = torch.device("cpu")
@@ -39,18 +38,18 @@ def predict(text, model_id, tokenizer_id):
     with torch.no_grad():
         logits = model(**inputs).logits
 
-    probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
-    output_pred = {model.config.id2label[i]: probs[i] for i in np.argsort(probs)[::-1]}
+    NUMS_DICT = {i: key for i, key in enumerate(sorted(EMOTION9_LABEL_NAMES.keys()))}
+    output_pred = {f"[{NUMS_DICT[i]}] {EMOTION9_LABEL_NAMES[NUMS_DICT[i]]}": probs[i] for i in np.argsort(probs)[::-1]}
     output_info = f'<p style="text-align: center; display: block">Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.</p>'
     return output_pred, output_info
 
-def predict_cap(text, language, domain):
+def predict_e6(text, language, domain):
     model_id = build_huggingface_path(language)
     tokenizer_id = "xlm-roberta-large"
     return predict(text, model_id, tokenizer_id)
 
 demo = gr.Interface(
-    fn=predict_cap,
+    fn=predict_e6,
     inputs=[gr.Textbox(lines=6, label="Input"),
             gr.Dropdown(languages, label="Language"),
             gr.Dropdown(domains.keys(), label="Domain")],