import transformers 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)) 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)