import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")

def chat(message, history):
    # Encode the conversation
    new_user_input = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")
    
    # Generate response
    bot_output = model.generate(
        new_user_input,
        max_length=100,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        temperature=0.8,
        top_p=0.9
    )
    
    # Decode and clean
    response = tokenizer.decode(bot_output[0], skip_special_tokens=True)
    
    # Remove the original message from response
    if response.startswith(message):
        response = response[len(message):].strip()
    
    # If still empty, return a default
    if not response:
        response = "Hello! How can I help you today?"
    
    return response

gr.ChatInterface(fn=chat, title="My Chatbot").launch()