dennis-fast commited on
Commit
77a444f
1 Parent(s): 3b5f022

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -0
app.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import transformers
2
+ import gradio as gr
3
+ import torch
4
+
5
+ from transformers import GPT2LMHeadModel, GPT2Tokenizer
6
+
7
+ tokenizer = GPT2Tokenizer.from_pretrained('microsoft/DialoGPT-small')
8
+ model = GPT2LMHeadModel.from_pretrained('microsoft/DialoGPT-small')
9
+
10
+ def chat(message, token_response):
11
+ token_message = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
12
+ token_response = model.generate(token_message, max_length=1000, pad_token_id=tokenizer.eos_token_id)
13
+ response = tokenizer.decode(token_response[:, token_message.shape[-1]:][0], skip_special_tokens=True)
14
+ return response, token_response
15
+
16
+ input = gr.inputs.Textbox(lines=2, label='User:')
17
+ output = gr.outputs.Textbox(label='Bot:')
18
+
19
+ gr.Interface(fn=chat,
20
+ title="DialoGPT-small",
21
+ inputs=[input, "state"],
22
+ outputs=[output, "state"],
23
+ allow_screenshot=False,
24
+ allow_flagging='never').launch()