Add application file

app.py

@@ -3,59 +3,22 @@ import gradio as gr
 import torch
 
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
 tokenizer = GPT2Tokenizer.from_pretrained('microsoft/DialoGPT-small')
 model = GPT2LMHeadModel.from_pretrained('microsoft/DialoGPT-small')
-model.eval()
-
-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))
+def chat(message, token_response):
+    token_message = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
+    token_response = model.generate(token_message, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(token_response[:, token_message.shape[-1]:][0], skip_special_tokens=True)
+    return response, token_response
 
-  return history, history
+input = gr.inputs.Textbox(lines=2, label='User:')
+output = gr.outputs.Textbox(label='Bot:')
 
-title = "DialoGPT-small"
-description = "Gradio demo for dialog using DialoGPT"
-iface = gr.Interface(
-    chat,
-    ["text", "state"],
-    ["chatbot", "state"],
+gr.Interface(fn=chat,
+    title="DialoGPT-small",
+    inputs=[input, "state"],
+    outputs=[output, "state"],
     allow_screenshot=False,
-    allow_flagging="never",
-    title=title,
-    description=description
-)
-iface.launch(debug=True)
+    allow_flagging='never').launch()