update chat state history, add initial greeting

README.md

@@ -1,10 +1,10 @@
 ---
 title: PsyPlus
 emoji: 🤖
-colorFrom: green
-colorTo: purple
+colorFrom: pink
+colorTo: blue
 sdk: gradio
-sdk_version: 3.11.0
+sdk_version: 3.10.0
 app_file: app.py
 pinned: false
 license: gpl-3.0

app.py

@@ -4,6 +4,7 @@ import matplotlib.pyplot as plt
 from threading import Timer
 import gradio as gr
 
+import torch
 from transformers import (
   GPT2LMHeadModel, GPT2Tokenizer,
   AutoModelForSequenceClassification, AutoTokenizer,
@@ -11,6 +12,8 @@ from transformers import (
 )
 # reference: https://huggingface.co/spaces/bentrevett/emotion-prediction
 # and https://huggingface.co/spaces/tareknaous/Empathetic-DialoGPT
+# 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/
 
 def euc_100():
   # 1,2,3. asks about the user's emotions and store data
@@ -77,16 +80,14 @@ def euc_100():
 
 
 def load_neural_emotion_detector():
-  model_name = "joeddav/distilbert-base-uncased-go-emotions-student"
+  model_name = 'joeddav/distilbert-base-uncased-go-emotions-student'
   model = AutoModelForSequenceClassification.from_pretrained(model_name)
+  model.eval()
   tokenizer = AutoTokenizer.from_pretrained(model_name)
   pipe = pipeline('text-classification', model=model, tokenizer=tokenizer, 
                   return_all_scores=True, truncation=True)
   return pipe
 
-def sort_predictions(predictions):
-  return sorted(predictions, key=lambda x: x['score'], reverse=True)
-
 def plot_emotion_distribution(predictions):
   fig, ax = plt.subplots()
   ax.bar(x=[i for i, _ in enumerate(prediction)],
@@ -98,9 +99,9 @@ def plot_emotion_distribution(predictions):
 
 def rulebase(text):
   keywords = {
-    'life_safety': ["death", "suicide", "murder", "to perish together", "jump off the building"],
-    'immediacy': ["now", "immediately", "tomorrow", "today"],
-    'manifestation': ["never stop", "every moment", "strong", "very"]
+    'life_safety': ['death', 'suicide', 'murder', 'to perish together', 'jump off the building'],
+    'immediacy': ['now', 'immediately', 'tomorrow', 'today'],
+    'manifestation': ['never stop', 'every moment', 'strong', 'very']
   }
 
   # if found dangerous kw/topics
@@ -127,7 +128,7 @@ def euc_200(text, testing=True):
   if not testing:
     pipe = load_neural_emotion_detector()
     prediction = pipe(text)[0]
-    prediction = sort_predictions(prediction)
+    prediction = sorted(predictions, key=lambda x: x['score'], reverse=True)
     plot_emotion_distribution(prediction)
   
   # get the most probable emotion. TODO: modify this part, may take sum of prob. over all negative emotion
@@ -174,46 +175,58 @@ def euc_200(text, testing=True):
   pass
 
 
-tokenizer = GPT2Tokenizer.from_pretrained("tareknaous/dialogpt-empathetic-dialogues")
-model = GPT2LMHeadModel.from_pretrained("tareknaous/dialogpt-empathetic-dialogues")
-model.eval()
-
-def chat(message, history):
-  history = history or []
+def _chat(message, history, model, tokenizer, args):
   eos = tokenizer.eos_token
-  input_str = eos.join([x for turn in history for x in turn] + [message])
+  history = history or {
+    'text': args.greeting, 
+    'input_ids': tokenizer.encode(args.greeting[-1][1] + eos, return_tensors='pt'), 
+  }
+  # TODO: only take the latest X turns, otherwise the text becomes longer and takes more time to process
 
-  bot_input_ids = tokenizer.encode(input_str, return_tensors='pt')
-  bot_output_ids = model.generate(bot_input_ids,
+  message_ids = tokenizer.encode(message + eos, return_tensors='pt')
+  history['input_ids'] = torch.cat([history['input_ids'], message_ids], dim=-1)
+
+  bot_output_ids = model.generate(history['input_ids'],
                                   max_length=1000,
                                   do_sample=True, top_p=0.9, temperature=0.8,
                                   pad_token_id=tokenizer.eos_token_id)
-  response = tokenizer.decode(bot_output_ids[:, bot_input_ids.shape[-1]:][0], 
+  response = tokenizer.decode(bot_output_ids[:, history['input_ids'].shape[-1]:][0], 
                              skip_special_tokens=True)
+  if args.run_on_own_server == 1: 
+    print((message, response), bot_output_ids[0][-10:])
 
-  history.append((message, response))
-  return history, history
+  history['input_ids'] = bot_output_ids
+  history['text'].append((message, response))
+  return history['text'], history
 
 
 if __name__ == '__main__':
   # euc_100()
   # euc_200('I am happy about my academic record.')
   parser = argparse.ArgumentParser()
-  parser.add_argument('--run_share_mode', type=int, default=0, help='if test on own server, need to use share mode')
+  parser.add_argument('--run_on_own_server', type=int, default=0, help='if test on own server, need to use share mode')
   args = parser.parse_args()
+  args.greeting = [('','Hi you!')]
+
+  tokenizer = GPT2Tokenizer.from_pretrained('tareknaous/dialogpt-empathetic-dialogues')
+  model = GPT2LMHeadModel.from_pretrained('tareknaous/dialogpt-empathetic-dialogues')
+  model.eval()
+  def chat(message, history):
+    return _chat(message, history, model, tokenizer, args)
 
-  title = "PsyPlus Empathetic Chatbot"
-  description = "Gradio demo for product of PsyPlus. Based on rule-based CBT and conversational AI model DialoGPT"
+  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=args.greeting)
   iface = gr.Interface(
     chat,
-    ["text", "state"],
-    ["chatbot", "state"],
-    allow_screenshot=False,
-    allow_flagging="never",
+    ['text', 'state'],
+    [chatbot, 'state'],
+    # css=".gradio-container {background-color: white}",
+    allow_flagging='never',
     title=title,
     description=description,
   )
-  if args.run_share_mode == 0:
+  if args.run_on_own_server == 0:
     iface.launch(debug=True)
   else:
-    iface.launch(debug=True, server_name="0.0.0.0", server_port=2022, share=True)
+    iface.launch(debug=True, server_name='0.0.0.0', server_port=2022, share=True)