Update app.py

app.py

@@ -73,6 +73,24 @@ def run_generate(text, bad_words):
     yo.append(e)
   return yo
   
+def BestProbs5(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(number_of_outputs)
+  best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
+  for i in best_words[0:number_of_outputs]:
+    #print(i)
+    print("\n")
+    g = (prompt + i)
+    st.write(g)
+    l = run_generate(g, "hey")
+    st.write(l)
+  
 def run_generate2(text, bad_words):
   yo = []
   input_ids = tokenizer2.encode(text, return_tensors='pt')
@@ -107,6 +125,8 @@ with st.form(key='my_form'):
     submit_button3 = st.form_submit_button(label='Submit Other Model')
     submit_button4 = st.form_submit_button(label='Submit Log Probs Other Model')
     
+    submit_button5 = st.form_submit_button(label='Most Prob')
+    
     if submit_button:
       translated_text = run_generate(text, bad_words)
       st.write(translated_text if translated_text else "No translation found")
@@ -137,4 +157,6 @@ with st.form(key='my_form'):
         probabilities = torch.nn.functional.softmax(logits)
         best_logits, best_indices = logits.topk(logs_outputs)
         best_words = [tokenizer2.decode([idx.item()]) for idx in best_indices]      
-        st.write(best_words)
+        st.write(best_words)
+    if submit_button5:
+      BestProbs5(prompt)