Updated inference.py

apps/inference.py

@@ -5,21 +5,25 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM
 from transformers import pipeline
 import os
 import json
+import random
+
+with open("config.json") as f:
+    cfg = json.loads(f.read())
 
 
 @st.cache(show_spinner=False, persist=True)
 def load_model(masked_text, model_name):
 
-    from_flax = False
-    if model_name == "flax-community/roberta-hindi":
-        from_flax = True
-    model = AutoModelForMaskedLM.from_pretrained(model_name, from_flax=from_flax)
+    model = AutoModelForMaskedLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+    nlp = pipeline("fill-mask", model=model, tokenizer=tokenizer)
+
     MASK_TOKEN = tokenizer.mask_token
+
     masked_text = masked_text.replace("<mask>", MASK_TOKEN)
-    nlp = pipeline("fill-mask", model=model, tokenizer=tokenizer)
     result_sentence = nlp(masked_text)
-    return result_sentence
+
+    return result_sentence[0]["sequence"], result_sentence[0]["token_str"]
 
 
 def app():
@@ -28,7 +32,7 @@ def app():
         unsafe_allow_html=True,
     )
     st.markdown(
-        "This demo uses pretrained RoBERTa variants for Mask Language Modelling (MLM)"
+        "This demo uses multiple hindi transformer models for Masked Language Modelling (MLM)."
     )
 
     target_text_path = "./mlm_custom/mlm_targeted_text.csv"
@@ -51,6 +55,56 @@ def app():
         ],
         ["flax-community/roberta-hindi"],
     )
+
+    models_list = list(cfg["models"].keys())
+
+        models = st.multiselect(
+        "Choose models",
+        models_list,
+        models_list[0],
+    )
+
+    target_text_path = "./mlm_custom/mlm_targeted_text.csv"
+    target_text_df = pd.read_csv(target_text_path)
+
+    texts = target_text_df["text"]
+
+    st.sidebar.title("Hindi MLM")
+
+    pick_random = st.sidebar.checkbox("Pick any random text")
+
+    results_df = pd.DataFrame(columns=["Model Name", "Filled Token", "Filled Text"])
+
+    model_names = []
+    filled_masked_texts = []
+    filled_tokens = []
+
+    if pick_random:
+        random_text = texts[random.randint(0, texts.shape[0] - 1)]
+        masked_text = st.text_area("Please type a masked sentence to fill", random_text)
+    else:
+        select_text = st.sidebar.selectbox("Select any of the following text", texts)
+        masked_text = st.text_area("Please type a masked sentence to fill", select_text)
+
+    # pd.set_option('max_colwidth',30)
+    if st.button("Fill the Mask!"):
+        with st.spinner("Filling the Mask..."):
+
+            for selected_model in models:
+
+                filled_sentence, filled_token = load_model(masked_text, cfg["models"][selected_model])
+                model_names.append(selected_model)
+                filled_tokens.append(filled_token)
+                filled_masked_texts.append(filled_sentence)
+
+            results_df["Model Name"] = model_names
+            results_df["Filled Token"] = filled_tokens
+            results_df["Filled Text"] = filled_masked_texts
+
+            st.table(results_df)
+
+
+
     if st.button("Fill the Mask!"):
         with st.spinner("Filling the Mask..."):
             models = []