7Vivek/Next-Word-Prediction-Streamlit

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