Update app.py

Adding lid models trained with Vuk data

app.py

@@ -155,14 +155,21 @@ def load_model_pipeline(model_name, file_name):
     
 
 
-model_1 = load_model(constants.MODEL_NAME, "model_v1.bin")
-model_2 = load_model(constants.MODEL_NAME, "model_v2.bin")
-model_3 = load_model(constants.MODEL_NAME, "model_v3.bin")
-openlid = load_model('laurievb/OpenLID', "model.bin")
-nllb = load_model('facebook/fasttext-language-identification', "model.bin")
+# model_1 = load_model(constants.MODEL_NAME, "model_v1.bin")
+# model_2 = load_model(constants.MODEL_NAME, "model_v2.bin")
+# model_3 = load_model(constants.MODEL_NAME, "model_v3.bin")
+# openlid = load_model('laurievb/OpenLID', "model.bin")
+# nllb = load_model('facebook/fasttext-language-identification', "model.bin")
+
+
+# MODELS
+model_xlmr_large = load_model_pipeline('dsfsi/za-xlmrlarge-lid', "model.bin")
+model_serengeti = load_model_pipeline('dsfsi/za-serengeti-lid', "model.bin")
+model_afriberta = load_model_pipeline('dsfsi/za-afriberta-lid', "model.bin")
+model_afroxlmr_base = load_model_pipeline('dsfsi/za-afro-xlmr-base-lid', "model.bin")
+model_afrolm        = load_model_pipeline('dsfsi/za-afrolm-lid', "model.bin")
 za_lid = load_model_pipeline('dsfsi/za-lid-bert', "model.bin")
 
-
 # @st.cache_resource
 def plot(label, prob):
 
@@ -197,16 +204,16 @@ def compute(sentences, version = 'v3'):
     """
     progress_text = "Computing Language..."
 
-    if version == 'nllb-218':
-        model_choice = nllb
-    elif version == 'openlid-201':
-        model_choice = openlid
-    elif version == 'v3':
-        model_choice = model_3   
-    elif version == 'v2':
-        model_choice = model_2
-    elif version=='v1':
-        model_choice = model_1
+    if version == 'xlmrlarge':
+        model_choice = model_xlmr_large
+    elif version == 'serengeti':
+        model_choice = model_serengeti
+    elif version == 'afriberta':
+        model_choice = model_afriberta   
+    elif version == 'afroxlmrbase':
+        model_choice = model_afroxlmr_base
+    elif version=='afrolm':
+        model_choice = model_afrolm
     else:
         model_choice = za_lid
         
@@ -270,8 +277,8 @@ with tab1:
 
     version = st.radio(
         "Choose model",
-        ["nllb-218", "openlid-201", "v1", "v2", "v3", "Za_LID"],
-        captions=["NLLB", "OpenLID", "GlotLID version 1", "GlotLID version 2", "GlotLID version 3 (More languages, better quality data)", "za BERT"],
+        ["xlmrlarge", "serengeti", "afriberta", "afroxlmrbase", "afrolm", "Za_LID"],
+        captions=["za-XLMR-Large", "za-Serengeti", "za-AfriBERTa", "za-Afro-XLMR-BASE", "za-AfroLM", "za-BERT"],
         index = 4,
         key = 'version_tab1',
         horizontal = True
@@ -309,8 +316,8 @@ with tab2:
 
     version = st.radio(
         "Choose model",
-        ["nllb-218", "openlid-201", "v1", "v2", "v3", "Za_LID"],
-        captions=["NLLB", "OpenLID", "GlotLID version 1", "GlotLID version 2 (more data and languages), GlotLID version 3 (More languages, better quality data)", "ZA BERT"],
+        ["xlmrlarge", "serengeti", "afriberta", "afroxlmrbase", "afrolm", "Za_LID"],
+        captions=["za-XLMR-Large", "za-Serengeti", "za-AfriBERTa", "za-Afro-XLMR-BASE", "za-AfroLM", "za-BERT"],
         index = 4,
         key = 'version_tab2',
         horizontal = True