add langcodes package

meedan · Aug 24, 2023 · 7b50c89 · 7b50c89
1 parent 764df76
commit 7b50c89
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/lib/model/fasttext.py b/lib/model/fasttext.py
@@ -3,6 +3,8 @@
 import fasttext
 from huggingface_hub import hf_hub_download
 
+from langcodes import *
+
 from lib.model.model import Model
 from lib import schemas
 
@@ -25,7 +27,7 @@ def respond(self, docs: Union[List[schemas.Message], schemas.Message]) -> List[s
         detectable_texts = [e.body.text for e in docs]
         detected_langs = []
         for text in detectable_texts:
-            detected_langs.append(self.model.predict(text)[0][0])
+            detected_langs.append(standardize_tag(self.model.predict(text)[0][0][9:], macro = True))
 
         for doc, detected_lang in zip(docs, detected_langs):
             doc.response = detected_lang