Added a Play again button

app.py

@@ -19,9 +19,6 @@ def get_models():
 
   st.write("The assistant is loaded and ready to use!")
   return model, tokenizer
-  
-  # st.session_state.model['model'] = model
-  # st.session_state.model['tokenizer'] = tokenizer
 
 model, tokenizer = get_models()
 
@@ -54,7 +51,7 @@ def return_top_k(sentence, k=10):
     probabilities = [round(float(prob), 2) for prob in decoded_probabilities]
     
     for pred in predictions:
-      if len(pred) < 2:
+      if (len(pred) < 2) | (pred in sentence.split()):
         predictions.pop(predictions.index(pred))  
      
   return predictions[:10]
@@ -78,13 +75,6 @@ if 'is_helpful' not in st.session_state:
 if 'descriptions' not in st.session_state:
   st.session_state.descriptions = []
   
-# if 'model' not in st.session_state:
-#   st.session_state.model = {'model': False, 'tokenizer': False}
-
-# if st.session_state.model['model'] == False:
-#   get_models()
-#   st.write("The assistant is loaded and ready to use!")
-  
 st.title("You name it!")
 
 with st.chat_message('user', avatar='simon.jpg'):
@@ -180,12 +170,19 @@ if st.session_state.actions[-1] == 'cue':
 
     elif b4:
       write_bot("Yay! I am happy I could be of help!")
+      new = st.button('Play again', key=63)
+      if new:
+        write_bot("Please describe your word!")
       guessed = True
+    
       break
     
     elif b5:
       write_bot("I am sorry I couldn't help you this time. See you soon!")
     st.session_state.actions.append('cue')
+    new = st.button('Play again', key=64)
+    if new:
+      write_bot("Please describe your word!")
     break
   
 #display user message in chat message container
@@ -195,7 +192,7 @@ if prompt:
     st.markdown(prompt)
   #add to history
   st.session_state.messages.append({'role': 'user', 'content': prompt})
-  yes = ['yes', 'again', 'Yes', 'sure', 'new word', 'yes!']
+  yes = ['yes', 'again', 'Yes', 'sure', 'new word', 'yes!', 'yep', 'yeah']
   if prompt in yes:
     write_bot("Please describe your word!")
   elif prompt == 'It is similar to the best place on earth':