fix global vs local chatbot

README.md

@@ -18,6 +18,7 @@ For more information about this product, please visit this notion [page](https:/
 
 ### 2022/12/20
 
+- DONE turning the chatbot to session varible so that different sessions will show different conversation
 - Chat flow will trigger euc 200 when detect a negative emotion with prob > threshold. Thus, only euc 100 and free chat consist of chat loop, while euc 200 will pop up sometimes. I set the trigger to NOT be regularly (currently one trigger once during the conversation), because trigger to much will bother users
 - Already fix the problem with dialog model. Now it's configured as the same as what it should be. Of course, that does not guarantee of good response
 - TODO is written in the main file already

app.py

@@ -8,9 +8,10 @@ reference:
 
 gradio vs streamlit 
   https://trojrobert.github.io/a-guide-for-deploying-and-serving-machine-learning-with-model-streamlit-vs-gradio/
-  https://gradio.app/interface_state/
+  https://gradio.app/interface_state/ -> global and local varible affect the separation of sessions
 
 TODO
+  Add command to reset/jump to a function, e.g >reset, >euc_100
   Add diagram in Gradio Interface showing sentimate analysis
   Gradio input timeout: cannot find a tutorial in Google -> don't know how to implement
   Personalize: create database, load and save data
@@ -40,8 +41,21 @@ def option():
   args = parser.parse_args()
   return args
 
+args = option()
+
+
+# store the list of messages that are showed in therapies and models as global variables
+# let all chat-session-wise variables placed in TherapyChatBot
+class ChatHelper:
+  # chat and emotion-detection models
+  ed_pipe = pipeline('text-classification', model=args.emotion_model, top_k=5, truncation=True)
+  ed_threshold = 0.3
+  dialog_model = GPT2LMHeadModel.from_pretrained(args.dialog_model)
+  dialog_tokenizer = GPT2Tokenizer.from_pretrained(args.dialog_model)
+  eos = dialog_tokenizer.eos_token
+  # tokenizer.__call__ -> input_ids, attention_mask
+  # tokenizer.encode -> only inputs_ids, which is required by model.generate function
 
-class ChatHelper: # store the list of messages that are showed in therapies
   invalid_input = 'Invalid input, my friend :) Plz input again'
   good_mood_over = 'Whether your good mood is over? Any other details that you would like to recall?'
   good_case = 'Nice to hear that!'
@@ -130,28 +144,19 @@ class TherapyChatBot:
     self.euc_100_emotion_degree = []
     self.already_trigger_euc_200 = False
 
-    # chat and emotion-detection models
-    self.ed_pipe = pipeline('text-classification', model=args.emotion_model, top_k=5, truncation=True)
-    self.ed_threshold = 0.3
-    self.dialog_model = GPT2LMHeadModel.from_pretrained(args.dialog_model)
-    self.dialog_tokenizer = GPT2Tokenizer.from_pretrained(args.dialog_model)
-    self.eos = self.dialog_tokenizer.eos_token
-    # tokenizer.__call__ -> input_ids, attention_mask
-    # tokenizer.encode -> only inputs_ids, which is required by model.generate function
-
     # chat history.
     # TODO: if we want to personalize and save the conversation,
     # we can load data from database
-    self.greeting = ChatHelper.greeting_template[self.chat_state]
-    self.history = {'input_ids': torch.tensor([[self.dialog_tokenizer.bos_token_id]]),
-                    'text': [('', self.greeting)]} if not self.account else open(f'database/{hash(self.account)}', 'rb')
+    self.greeting = [('', ChatHelper.greeting_template[self.chat_state])]
+    self.history = {'input_ids': torch.tensor([[ChatHelper.dialog_tokenizer.bos_token_id]]),
+                    'text': self.greeting} if not self.account else open(f'database/{hash(self.account)}', 'rb')
     if 'euc_100' in self.chat_state:
       self.chat_state = 'euc_100.q.0'
 
   def __call__(self, message, prefix=''):
     # if prefix != None, which means this function is called from euc_200, thus already detected the negative emotion
     if (not prefix) and self.chat_state != 'euc_200' and not self.already_trigger_euc_200:
-      prediction = self.ed_pipe(message)[0]
+      prediction = ChatHelper.ed_pipe(message)[0]
       prediction = sorted(prediction, key=lambda x: x['score'], reverse=True)
       if self.run_on_own_server:
         print(prediction)
@@ -160,7 +165,7 @@ class TherapyChatBot:
 
     # if message is negative, change state immediately
     if ((not prefix) and self.chat_state != 'euc_200' and not self.already_trigger_euc_200) and \
-      (emotion['label'] in ChatHelper.negative_emotions and emotion['score'] > self.ed_threshold):
+      (emotion['label'] in ChatHelper.negative_emotions and emotion['score'] > ChatHelper.ed_threshold):
       self.chat_state_prev = self.chat_state
       self.chat_state = 'euc_200'
       self.message_prev = message
@@ -171,7 +176,7 @@ class TherapyChatBot:
     elif self.chat_state.startswith('euc_100'):
       response = self.euc_100(message)
       if self.chat_state == 'free_chat':
-        last_two_turns_ids = self.dialog_tokenizer.encode(message + self.eos, return_tensors='pt')
+        last_two_turns_ids = ChatHelper.dialog_tokenizer.encode(message + ChatHelper.eos, return_tensors='pt')
         self.history['input_ids'] = torch.cat([self.history['input_ids'], last_two_turns_ids], dim=-1)
 
     elif self.chat_state.startswith('euc_200'):
@@ -185,7 +190,6 @@ class TherapyChatBot:
       self.history['text'].append((self.message_prev, response))
     else:
       self.history['text'].append((message, response))
-    return self.history['text']
 
   def euc_100(self, x):
     _, subsection, entry = self.chat_state.split('.')
@@ -251,23 +255,23 @@ class TherapyChatBot:
     message = self.message_prev
     self.message_prev = x
     self.chat_state = self.chat_state_prev
-    return self.__call__(message, response)
+    return self.__call__(message, prefix=response)
 
   def free_chat(self, message):
-    message_ids = self.dialog_tokenizer.encode(message + self.eos, return_tensors='pt')
+    message_ids = ChatHelper.dialog_tokenizer.encode(message + ChatHelper.eos, return_tensors='pt')
     self.history['input_ids'] = torch.cat([self.history['input_ids'], message_ids], dim=-1)
     input_ids = self.history['input_ids'].clone()
 
     while True:
-      bot_output_ids = self.dialog_model.generate(input_ids, max_length=1000,
+      bot_output_ids = ChatHelper.dialog_model.generate(input_ids, max_length=1000,
                                           do_sample=True, top_p=0.9, temperature=0.8, num_beams=2,
-                                          pad_token_id=self.dialog_tokenizer.eos_token_id)
-      response = self.dialog_tokenizer.decode(bot_output_ids[0][input_ids.shape[-1]:],
+                                          pad_token_id=ChatHelper.dialog_tokenizer.eos_token_id)
+      response = ChatHelper.dialog_tokenizer.decode(bot_output_ids[0][input_ids.shape[-1]:],
                                        skip_special_tokens=True)
       if response.strip() != '':
         break
-      elif input_ids[0].tolist().count(self.dialog_tokenizer.eos_token_id) > 0:
-        idx = input_ids[0].tolist().index(self.dialog_tokenizer.eos_token_id)
+      elif input_ids[0].tolist().count(ChatHelper.dialog_tokenizer.eos_token_id) > 0:
+        idx = input_ids[0].tolist().index(ChatHelper.dialog_tokenizer.eos_token_id)
         input_ids = input_ids[:, (idx+1):]
       else:
         input_ids = message_ids
@@ -282,20 +286,22 @@ class TherapyChatBot:
     return response
 
 
-if __name__ == '__main__':
-  args = option()
-  chat = TherapyChatBot(args)
+if __name__ == '__main__':  
+  def chat(message, bot):
+    bot = bot or TherapyChatBot(args)
+    bot(message)
+    return bot.history['text'], bot
 
   title = 'PsyPlus Empathetic Chatbot'
   description = 'Gradio demo for product of PsyPlus. Based on rule-based CBT and conversational AI model DialoGPT'
-  chatbot = gr.Chatbot(value=chat.history['text'])
+  greeting = [('', ChatHelper.greeting_template[args.initial_chat_state])]
+  chatbot = gr.Chatbot(value=greeting)
   iface = gr.Interface(
-    chat, 'text', chatbot,
+    chat, ['text', 'state'], [chatbot, 'state'],
     allow_flagging='never', title=title, description=description,
   )
 
-  # iface.queue(concurrency_count=5)
   if args.run_on_own_server == 0:
     iface.launch(debug=True)
   else:
-    iface.launch(debug=True, share=True) # server_name='0.0.0.0', server_port=2022
+    iface.launch(debug=True, share=True)