Update app.py

app.py

@@ -1,43 +1,44 @@
-from transformers import GPTNeoForCausalLM, GPT2Tokenizer
-import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+
+def predict(input, history=[]):
+    # tokenize the new input sentence
+    new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors='pt')
+
+    # append the new user input tokens to the chat history
+    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
+
+    # generate a response 
+    history = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id).tolist()
+
+    # convert the tokens to text, and then split the responses into lines
+    response = tokenizer.decode(history[0]).split("<|endoftext|>")
+    response.remove("")
+    
+    # write some HTML
+    html = "<div class='chatbot'>"
+    for m, msg in enumerate(response):
+        cls = "user" if m%2 == 0 else "bot"
+        html += "<div class='msg {}'> {}</div>".format(cls, msg)
+    html += "</div>"
+    
+    return html, history
 
-model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
-tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
-
-prompt = """This is a discussion between a person and Hassan Kane, an entrepreneur. 
-person: What are you working on?
-Hassan: This new AI community building the future of Africa
-person: Where are you?
-Hassan: In Lagos for a week, then Paris or London.
-person: How's it going? 
-Hassan: Not bad.. Just trying to hit EV (escape velocity) with my startup
-person: """
-
-def my_split(s, seps):
-    res = [s]
-    for sep in seps:
-        s, res = res, []
-        for seq in s:
-            res += seq.split(sep)
-    return res
-
-# input = "Who are you?"
-def chat_base(input):
-  p = prompt + input
-  input_ids = tokenizer(p, return_tensors="pt").input_ids
-  gen_tokens = model.generate(input_ids, do_sample=True, temperature=0.7, max_length=150,)
-  gen_text = tokenizer.batch_decode(gen_tokens)[0]
-  # print(gen_text)
-  result = gen_text[len(p):]   
-  # print(">", result)
-  result = my_split(result, [']', '\n'])[1]
-  # print(">>", result)
-  if "Hassan: " in result:
-   result = result.split("Hassan: ")[-1]
-# print(">>>", result)
-  return result
-  
 import gradio as gr
 
-iface = gr.Interface(chat_base, gr.inputs.Textbox(label="Ask Hassan a Question"), "text", allow_screenshot=False, allow_flagging=False)
-iface.launch()
+css = """
+.chatbox {display:flex;flex-direction:column}
+.msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%}
+.msg.user {background-color:cornflowerblue;color:white}
+.msg.bot {background-color:lightgray;align-self:self-end}
+.footer {display:none !important}
+"""
+
+gr.Interface(fn=predict,
+             theme="default",
+             inputs=[gr.inputs.Textbox(placeholder="How are you?"), "state"],
+             outputs=["html", "state"],
+             css=css).launch()