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()