Update app.py

app.py

@@ -166,12 +166,58 @@ def LogProbs(prompt):
   print(df)
   st.write(df)
   return df
+  
+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(5)
+  best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
+  for i in best_words[0:5]:
+    #print(i)
+    print("\n")
+    g = (prompt + i)
+    st.write(g)
+    l = run_generate(g, "hey")
+    st.write(l)
+    
+def run_generate(text, bad_words):
+  yo = []
+  input_ids = tokenizer.encode(text, return_tensors='pt')
+  res = len(tokenizer.encode(text))
+  bad_words = bad_words.split()
+  bad_word_ids = [[7829], [40940]]
+  for bad_word in bad_words: 
+    bad_word = " " + bad_word
+    ids = tokenizer(bad_word).input_ids
+    bad_word_ids.append(ids)
+  sample_outputs = model.generate(
+    input_ids,
+    do_sample=True, 
+    max_length= res + 5, 
+    min_length = res + 5, 
+    top_k=50,
+    temperature=1.0,
+    num_return_sequences=3,
+    bad_words_ids=bad_word_ids
+  )
+  for i in range(3):
+    e = tokenizer.decode(sample_outputs[i])
+    e = e.replace(text, "")
+    yo.append(e)
+  print(yo)
+  return yo
 
 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')
+    submit_button4 = st.form_submit_button(label='Get Top')
 
     if submit_button:
       with torch.no_grad():
@@ -198,4 +244,6 @@ with st.form(key='my_form'):
     if submit_button3:
         print("----")
         st.write("___")
-        st.write(BestProbs)
+        st.write(BestProbs)
+    if submit_button4:
+      BestProbs5(prompt)