Update app.py

app.py

@@ -81,7 +81,7 @@ def BestProbs(prompt):
   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_logits, best_indices = logits.topk(10)
   best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
   for i in best_words[0:10]:
     print("_______")
@@ -102,12 +102,14 @@ def BestProbs2(prompt):
   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]:
+    print(i)
     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')
+    submit_button3 = st.form_submit_button(label='Fast Forward 2.0')
 
     if submit_button:
       with torch.no_grad():
@@ -125,6 +127,11 @@ with st.form(key='my_form'):
         words = []              
         st.write(best_words)
       if submit_button2:
+          print("----")
+          st.write("___")
+          m = BestProbs(prompt)
+          st.write(m)
+      if submit_button3:
+          print("----")
           st.write("___")
-          ok = BestProbs(prompt)
-          st.write(ok)
+          st.write(BestProbs)