Testing indicberta

apps/inference.py

@@ -13,21 +13,13 @@ def load_model(masked_text, model_name):
     from_flax = False
     if model_name == "flax-community/roberta-hindi":
         from_flax = True
-
-    if model_name == "neuralspace-reverie/indic-transformers-hi-bert":
-        masked_text = masked_text.replace("<mask>", "[MASK]")
-
-    elif model_name == "ai4bharat/indic-bert":
-        masked_text = masked_text.replace("<mask>", "([MASK])")
-
-    st.write(model_name, masked_text)
-
+    # st.write(model_name, masked_text)
     model = AutoModelForMaskedLM.from_pretrained(model_name, from_flax=from_flax)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+    MASK_TOKEN = tokenizer.mask_token
+    masked_text = masked_text.replcae("<mask>", MASK_TOKEN)
     nlp = pipeline("fill-mask", model=model, tokenizer=tokenizer)
-
     result_sentence = nlp(masked_text)
-
     return result_sentence