| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium") | |
| model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium") | |
| chat_history_ids = None | |
| def chat(input_text, history=[]): | |
| global chat_history_ids | |
| new_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt') | |
| if chat_history_ids is not None: | |
| bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1) | |
| else: | |
| bot_input_ids = new_input_ids | |
| chat_history_ids = model.generate( | |
| bot_input_ids, | |
| max_length=1000, | |
| pad_token_id=tokenizer.eos_token_id, | |
| do_sample=True, | |
| top_k=50, | |
| top_p=0.95, | |
| temperature=0.75, | |
| ) | |
| response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True) | |
| history.append((input_text, response)) | |
| return history, history | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## 💖 AI Girlfriend Chat (DialoGPT)") | |
| chatbot = gr.Chatbot() | |
| msg = gr.Textbox(label="Type something...") | |
| state = gr.State([]) | |
| msg.submit(chat, [msg, state], [chatbot, state]) | |
| demo.launch() | |