updated inference.py

apps/inference.py

@@ -1,50 +1,58 @@
 from pandas.io.formats.format import return_docstring
 import streamlit as st
 import pandas as pd
-from transformers import AutoTokenizer,AutoModelForMaskedLM
+from transformers import AutoTokenizer, AutoModelForMaskedLM
 from transformers import pipeline
 import os
 import json
 
-@st.cache(show_spinner=False,persist=True)
-def load_model(masked_text,model_name):
+
+@st.cache(show_spinner=False, persist=True)
+def load_model(masked_text, model_name):
 
     model = AutoModelForMaskedLM.from_pretrained(model_name, from_flax=True)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    nlp = pipeline('fill-mask', model=model, tokenizer=tokenizer)
+    nlp = pipeline("fill-mask", model=model, tokenizer=tokenizer)
 
     result_sentence = nlp(masked_text)
 
-    return result_sentence[0]['sequence']
+    return result_sentence
+
 
 def app():
-    st.markdown("<h1 style='text-align: center; color: green;'>RoBERTa Hindi</h1>", unsafe_allow_html=True)
     st.markdown(
-    "This demo uses pretrained RoBERTa variants for Mask Language Modelling (MLM)"
+        "<h1 style='text-align: center; color: green;'>RoBERTa Hindi</h1>",
+        unsafe_allow_html=True,
+    )
+    st.markdown(
+        "This demo uses pretrained RoBERTa variants for Mask Language Modelling (MLM)"
     )
 
-    target_text_path = './mlm_custom/mlm_targeted_text.csv'
+    target_text_path = "./mlm_custom/mlm_targeted_text.csv"
     target_text_df = pd.read_csv(target_text_path)
-    
-    texts = target_text_df['text']
-    
+
+    texts = target_text_df["text"]
+
     st.markdown("""## Select any of the following text : """)
-    masked_text = st.selectbox('',
-     texts)
+    masked_text = st.selectbox("", texts)
 
-    st.write('You selected:', masked_text)
+    st.write("You selected:", masked_text)
 
     models = st.multiselect(
         "Choose models",
-        ['flax-community/roberta-hindi','mrm8488/HindiBERTa','ai4bharat/indic-bert',\
-        'neuralspace-reverie/indic-transformers-hi-bert', 
-          'surajp/RoBERTa-hindi-guj-san'],
-        ["flax-community/roberta-hindi"]
-    )   
+        [
+            "flax-community/roberta-hindi",
+            "mrm8488/HindiBERTa",
+            "ai4bharat/indic-bert",
+            "neuralspace-reverie/indic-transformers-hi-bert",
+            "surajp/RoBERTa-hindi-guj-san",
+        ],
+        ["flax-community/roberta-hindi"],
+    )
 
     selected_model = models[0]
 
-    if st.button('Fill the Mask!'):
-         with st.spinner("Filling the Mask..."):
-            filled_sentence = load_model(masked_text,selected_model)
+    if st.button("Fill the Mask!"):
+        with st.spinner("Filling the Mask..."):
+            filled_sentence = load_model(masked_text, selected_model)
             st.write(filled_sentence)