roberta-testing / apps /inference.py
prateekagrawal's picture
Updated requirements.txt
2fe741d
raw history blame
No virus
1.75 kB
from pandas.io.formats.format import return_docstring
import streamlit as st
import pandas as pd
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):
model = AutoModelForMaskedLM.from_pretrained(model_name, from_flax=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# tokenizer.save_pretrained('exported_pytorch_model')
# model.save_pretrained('exported_pytorch_model')
nlp = pipeline('fill-mask', model="exported_pytorch_model")
result_sentence = nlp(masked_text)
return result_sentence[0]['sequence']
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)"
)
target_text_path = './mlm_custom/mlm_targeted_text.csv'
target_text_df = pd.read_csv(target_text_path)
texts = target_text_df['text']
st.markdown("""## Select any of the following text : """)
masked_text = st.selectbox('',
texts)
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"]
)
selected_model = models[0]
if st.button('Fill the Mask!'):
with st.spinner("Filling the Mask..."):
filled_sentence = load_model(masked_text,selected_model)
st.write(filled_sentence)