kovacsvi commited on
Commit
0a077e0
·
1 Parent(s): 32d95f0

use different dicts for v1-v2 emotion9 predictions

Browse files
Files changed (1) hide show
  1. interfaces/emotion9.py +7 -2
interfaces/emotion9.py CHANGED
@@ -7,7 +7,7 @@ from transformers import AutoModelForSequenceClassification
7
  from transformers import AutoTokenizer
8
  from huggingface_hub import HfApi
9
 
10
- from label_dicts import EMOTION9_LABEL_NAMES
11
 
12
  from .utils import is_disk_full, release_model
13
 
@@ -20,10 +20,15 @@ domains = {
20
  "parliamentary speech": "parlspeech",
21
  }
22
 
 
 
 
 
 
23
  def build_huggingface_path(language: str):
24
  language = language.lower()
25
  base_model_id = f"poltextlab/xlm-roberta-large-pooled-{language}-emotions9"
26
- if language in ["slovak", "polish", "czech", "hungarian"]:
27
  return base_model_id + "-v2"
28
  return base_model_id
29
 
 
7
  from transformers import AutoTokenizer
8
  from huggingface_hub import HfApi
9
 
10
+ from label_dicts import EMOTION9_LABEL_NAMES, EMOTION9_V2_LABEL_NAMES
11
 
12
  from .utils import is_disk_full, release_model
13
 
 
20
  "parliamentary speech": "parlspeech",
21
  }
22
 
23
+ USE_V2 = False
24
+ if language in ["slovak", "polish", "czech", "hungarian"]:
25
+ EMOTION9_LABEL_NAMES = EMOTION9_V2_LABEL_NAMES
26
+ USE_V2 = True
27
+
28
  def build_huggingface_path(language: str):
29
  language = language.lower()
30
  base_model_id = f"poltextlab/xlm-roberta-large-pooled-{language}-emotions9"
31
+ if USE_V2:
32
  return base_model_id + "-v2"
33
  return base_model_id
34