Spaces:
Build error
Build error
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| # Load model and tokenizer | |
| model_name = "microsoft/DialoGPT-small" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| def respond(message, chat_history, chat_history_ids): | |
| # Encode user input | |
| new_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt") | |
| # Append to chat history | |
| if chat_history_ids is not None: | |
| input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1) | |
| else: | |
| input_ids = new_input_ids | |
| # Generate response | |
| chat_history_ids = model.generate( | |
| input_ids, | |
| max_length=1000, | |
| pad_token_id=tokenizer.eos_token_id, | |
| no_repeat_ngram_size=3, | |
| do_sample=True, | |
| top_k=50, | |
| top_p=0.95, | |
| temperature=0.8 | |
| ) | |
| # Decode response | |
| response = tokenizer.decode( | |
| chat_history_ids[:, input_ids.shape[-1]:][0], | |
| skip_special_tokens=True | |
| ) | |
| # Update conversation history | |
| chat_history.append((message, response)) | |
| return "", chat_history, chat_history_ids | |
| with gr.Blocks() as demo: | |
| # Store model's conversation history | |
| state = gr.State() | |
| gr.Markdown("## DialoGPT Chatbot") | |
| chatbot = gr.Chatbot() | |
| msg = gr.Textbox(label="Your Message") | |
| clear = gr.Button("Clear History") | |
| msg.submit( | |
| respond, | |
| [msg, chatbot, state], | |
| [msg, chatbot, state] | |
| ) | |
| clear.click(lambda: (None, None), outputs=[chatbot, state], queue=False) | |
| demo.launch() |