unmasker_replace.py

from transformers import pipeline
import random

unmasker = pipeline('fill-mask', model='bert-base-cased')

input_text = "I went to see a movie in the theater"

orig_text_list = input_text.split()
len_input = len(orig_text_list)
#Random index where we want to replace the word 
rand_idx = random.randint(1,len_input-1)
orig_word = orig_text_list[rand_idx]
new_text_list = orig_text_list.copy()
new_text_list[rand_idx] = '[MASK]'
new_mask_sent = ' '.join(new_text_list)
print("Masked sentence->",new_mask_sent)
#I went to [MASK] a movie in the theater

augmented_text_list = unmasker(new_mask_sent)
#To ensure new word and old word are not name
for res in augmented_text_list:
  if res['token_str'] != orig_word:
    augmented_text = res['sequence']
    break
print("Augmented text->",augmented_text)
#I went to watch a movie in the theater

Hi, thanks for pointing it out, while copying from my notebook to gist, I forgot to include that line - 'orig_word' is initial the word at the random index, so its - orig_word = orig_text_list[rand_idx]