Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| # Load model and tokenizer | |
| model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small") | |
| tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small") | |
| def chat_response(message, history): | |
| try: | |
| # Build conversation history string | |
| chat_history_ids = None | |
| for human_msg, bot_msg in history: | |
| # Encode user message | |
| user_input_ids = tokenizer.encode( | |
| human_msg + tokenizer.eos_token, | |
| return_tensors='pt' | |
| ) | |
| # Encode bot response | |
| bot_output_ids = tokenizer.encode( | |
| bot_msg + tokenizer.eos_token, | |
| return_tensors='pt' | |
| ) | |
| # Build full conversation | |
| if chat_history_ids is None: | |
| chat_history_ids = torch.cat([user_input_ids, bot_output_ids], dim=-1) | |
| else: | |
| chat_history_ids = torch.cat([chat_history_ids, user_input_ids, bot_output_ids], dim=-1) | |
| # Add new user message | |
| new_user_input_ids = tokenizer.encode( | |
| message + tokenizer.eos_token, | |
| return_tensors='pt' | |
| ) | |
| # Generate response | |
| chat_history_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if chat_history_ids is not None else new_user_input_ids | |
| # Generate bot response | |
| bot_output_ids = model.generate( | |
| chat_history_ids, | |
| max_length=1000, | |
| pad_token_id=tokenizer.eos_token_id, | |
| no_repeat_ngram_size=3, | |
| do_sample=True, | |
| top_k=100, | |
| top_p=0.7, | |
| temperature=0.8 | |
| ) | |
| # Extract only the bot's response (remove history) | |
| response = tokenizer.decode( | |
| bot_output_ids[:, chat_history_ids.shape[-1]:][0], | |
| skip_special_tokens=True | |
| ) | |
| return response | |
| except Exception as e: | |
| return f"Error: {str(e)}" | |
| # Create chat interface | |
| demo = gr.ChatInterface( | |
| chat_response, | |
| title="DialoGPT Chatbot", | |
| examples=["Hello!", "What's AI?", "Tell me a joke"], | |
| type="messages" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |