Update app.py

app.py

@@ -72,9 +72,40 @@ Translated into the Style of Abraham Lincoln: corn fields ( permeate illinois /
 
 informal english:"""
 
+
+def BestProbs(prompt):
+  prompt = prompt.strip()
+  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(200)
+  best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
+  for i in best_words[0:10]:
+    print("_______")
+    st.write(f"${i} $\n")
+    m = (prompt + f"{i}")
+    BestProbs2(m)
+
+def BestProbs2(prompt):
+  prompt = prompt.strip()
+  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(10)
+  best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
+  for i in best_words[0:10]:
+    st.write(i)
+
 with st.form(key='my_form'):
     prompt = st.text_area(label='Enter sentence', value=g)
     submit_button = st.form_submit_button(label='Submit')
+    submit_button2 = st.form_submit_button(label='Fast Forward')
 
     if submit_button:
       with torch.no_grad():
@@ -90,4 +121,6 @@ with st.form(key='my_form'):
         text.append(best_indices[0].item())
         best_probabilities = probabilities[best_indices].tolist()
         words = []              
-        st.write(best_words)
+        st.write(best_words)
+      if submit_button2:
+          BestProbs(prompt)