Spaces:
Runtime error
Runtime error
File size: 3,375 Bytes
df6828f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import os
import streamlit as st
import torch
import string
from transformers import BertTokenizer, BertForMaskedLM
st.set_page_config(page_title='Next Word Prediction Model', page_icon=None, layout='centered', initial_sidebar_state='auto')
@st.cache()
def load_model(model_name):
try:
if model_name.lower() == "bert":
bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
bert_model = BertForMaskedLM.from_pretrained('bert-base-uncased').eval()
return bert_tokenizer,bert_model
except Exception as e:
pass
#use joblib to fast your function
def decode(tokenizer, pred_idx, top_clean):
ignore_tokens = string.punctuation + '[PAD]'
tokens = []
for w in pred_idx:
token = ''.join(tokenizer.decode(w).split())
if token not in ignore_tokens:
tokens.append(token.replace('##', ''))
return '\n'.join(tokens[:top_clean])
def encode(tokenizer, text_sentence, add_special_tokens=True):
text_sentence = text_sentence.replace('<mask>', tokenizer.mask_token)
# if <mask> is the last token, append a "." so that models dont predict punctuation.
if tokenizer.mask_token == text_sentence.split()[-1]:
text_sentence += ' .'
input_ids = torch.tensor([tokenizer.encode(text_sentence, add_special_tokens=add_special_tokens)])
mask_idx = torch.where(input_ids == tokenizer.mask_token_id)[1].tolist()[0]
return input_ids, mask_idx
def get_all_predictions(text_sentence, top_clean=5):
# ========================= BERT =================================
input_ids, mask_idx = encode(bert_tokenizer, text_sentence)
with torch.no_grad():
predict = bert_model(input_ids)[0]
bert = decode(bert_tokenizer, predict[0, mask_idx, :].topk(top_k).indices.tolist(), top_clean)
return {'bert': bert}
def get_prediction_eos(input_text):
try:
input_text += ' <mask>'
res = get_all_predictions(input_text, top_clean=int(top_k))
return res
except Exception as error:
pass
try:
st.markdown("<h1 style='text-align: center;'>Next Word Prediction</h1>", unsafe_allow_html=True)
st.markdown("<h4 style='text-align: center; color: #B2BEB5;'><i>Keywords : BertTokenizer, BertForMaskedLM, Pytorch</i></h4>", unsafe_allow_html=True)
st.sidebar.text("Next Word Prediction Model")
top_k = st.sidebar.slider("Select How many words do you need", 1 , 25, 1) #some times it is possible to have less words
print(top_k)
model_name = st.sidebar.selectbox(label='Select Model to Apply', options=['BERT', 'XLNET'], index=0, key = "model_name")
bert_tokenizer, bert_model = load_model(model_name)
input_text = st.text_area("Enter your text here")
#click outside box of input text to get result
res = get_prediction_eos(input_text)
answer = []
print(res['bert'].split("\n"))
for i in res['bert'].split("\n"):
answer.append(i)
answer_as_string = " ".join(answer)
st.text_area("Predicted List is Here",answer_as_string,key="predicted_list")
st.image('https://freepngimg.com/download/keyboard/6-2-keyboard-png-file.png',use_column_width=True)
st.markdown("<h6 style='text-align: center; color: #808080;'>Created By <a href='https://github.com/7Vivek'>Vivek</a> - Checkout complete project <a href='https://github.com/7Vivek/Next-Word-Prediction-Streamlit'>here</a></h6>", unsafe_allow_html=True)
except Exception as e:
print("SOME PROBLEM OCCURED")
|