Update app.py

app.py

@@ -189,6 +189,24 @@ def BestProbs5(prompt):
     l = run_generate(g, "hey")
     st.write(l)
     
+def syn(prompt, word):
+  prompt = prompt.strip()
+  prompt = prompt.split(word)[0]
+  prompt = "Translated into the Style of Abraham Lincoln: " + prompt + "( " + word + " /"
+  #prompt = prompt.replace("/ ", "/")
+  print(prompt)
+  text = tokenizer.encode(prompt)
+  myinput, past_key_values = torch.tensor([text]), None
+  myinput = myinput
+  logits, past_key_values = model(myinput, past_key_values = past_key_values, return_dict=False)
+  logits = logits[0,-1]
+  probabilities = torch.nn.functional.softmax(logits)
+  best_logits, best_indices = logits.topk(100)
+  best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
+  for i in best_words[0:100]:
+    print(prompt)
+    st.write(i)
+    
 def run_generate(text, bad_words):
   yo = []
   input_ids = tokenizer.encode(text, return_tensors='pt')
@@ -222,6 +240,7 @@ with st.form(key='my_form'):
     submit_button2 = st.form_submit_button(label='Fast Forward')
     submit_button3 = st.form_submit_button(label='Fast Forward 2.0')
     submit_button4 = st.form_submit_button(label='Get Top')
+    submit_button5 = st.form_submit_button(label='Synonyms')
 
     if submit_button:
       with torch.no_grad():
@@ -250,4 +269,7 @@ with st.form(key='my_form'):
         st.write("___")
         st.write(BestProbs)
     if submit_button4:
-      BestProbs5(prompt)
+      BestProbs5(prompt)
+    if submit_button5:
+      bad_words = st.text_input("Word Or Phrase You Want Syn For", "yielded to")
+      syn(prompt, bad_words)