Update app.py

app.py

@@ -1,15 +1,50 @@
-from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration
-import gradio as grad
-
-chat_tkn = BlenderbotTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
-mdl = BlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-400M-distill")
-
-def createHistory(message):
-    history = grad.get_state() or []
+from transformers import GPTNeoForCausalLM, GPT2Tokenizer
+import gradio as gr
+
+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
+
+def chat(message):
+    history = gr.get_state() or []
     print(history)
-    response = chat(message)
+    response = chat_base(message)
     history.append((message, response))
-    grad.set_state(history)
+    gr.set_state(history)
     html = "<div class='chatbot'>"
     for user_msg, resp_msg in history:
         html += f"<div class='user_msg'>{user_msg}</div>"
@@ -17,18 +52,5 @@ def createHistory(message):
     html += "</div>"
     return response
 
-def chat(input):
-     
-    tkn_ids = chat_tkn(input+ chat_tkn.eos_token, return_tensors='pt')
-
-    # bot responds
-    chat_ids = mdl.generate(**tkn_ids)
-
-    # print bot response
-    response= "Alicia: {}".format(chat_tkn.decode(chat_ids[0], skip_special_tokens=True))
-
-    return response
-
-out=grad.Textbox(lines=20, label="dialog", placeholder="start conversation")
-grad.Interface(createHistory, inputs="text",outputs=out).launch()
-
+iface = gr.Interface(chat_base, gr.inputs.Textbox(label="Ask Hassan a Question"), "text", allow_screenshot=False, allow_flagging=False)
+iface.launch()