Add application file

app.py

@@ -1,7 +1,61 @@
+import transformers
 import gradio as gr
+import torch
 
-def greet(name):
-    return "Hello " + name + "!!"
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+tokenizer = GPT2Tokenizer.from_pretrained('microsoft/DialoGPT-small')
+model = GPT2LMHeadModel.from_pretrained('microsoft/DialoGPT-small')
+model.eval()
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+def chat(message, history):
+  history = history or []
+  new_user_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
+  
+  if len(history) > 0 and len(history) < 2:
+    for i in range(0,len(history)):
+      encoded_message = tokenizer.encode(history[i][0] + tokenizer.eos_token, return_tensors='pt')
+      encoded_response = tokenizer.encode(history[i][1] + tokenizer.eos_token, return_tensors='pt')
+      if i == 0:
+        chat_history_ids = encoded_message
+        chat_history_ids = torch.cat([chat_history_ids,encoded_response], dim=-1)
+      else:
+        chat_history_ids = torch.cat([chat_history_ids,encoded_message], dim=-1)
+        chat_history_ids = torch.cat([chat_history_ids,encoded_response], dim=-1)
+      
+      bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
+
+  elif len(history) >= 2:
+      for i in range(len(history)-1, len(history)):
+        encoded_message = tokenizer.encode(history[i][0] + tokenizer.eos_token, return_tensors='pt')
+        encoded_response = tokenizer.encode(history[i][1] + tokenizer.eos_token, return_tensors='pt')
+        if i == (len(history)-1):
+          chat_history_ids = encoded_message
+          chat_history_ids = torch.cat([chat_history_ids,encoded_response], dim=-1)
+        else:
+          chat_history_ids = torch.cat([chat_history_ids,encoded_message], dim=-1)
+          chat_history_ids = torch.cat([chat_history_ids,encoded_response], dim=-1)
+      
+      bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
+  
+  elif len(history) == 0:
+    bot_input_ids =  new_user_input_ids
+
+  chat_history_ids = model.generate(bot_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(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+    
+  history.append((message, response))
+
+  return history, history
+
+title = "DialoGPT-small"
+description = "Gradio demo for dialog using DialoGPT"
+iface = gr.Interface(
+    chat,
+    ["text", "state"],
+    ["chatbot", "state"],
+    allow_screenshot=False,
+    allow_flagging="never",
+    title=title,
+    description=description
+)
+iface.launch(debug=True)