import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL_NAME = "microsoft/DialoGPT-medium"
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
model.eval()
print("Model ready!")


def respond(message, history):
    print(f"=== respond called with message: {message!r} ===")
    try:
        text = str(message) + tokenizer.eos_token
        input_ids = tokenizer.encode(text, return_tensors="pt")

        with torch.no_grad():
            output_ids = model.generate(
                input_ids,
                max_new_tokens=50,
                pad_token_id=tokenizer.eos_token_id,
                do_sample=True,
                top_p=0.9,
                temperature=0.8,
            )

        reply = tokenizer.decode(
            output_ids[:, input_ids.shape[-1]:][0],
            skip_special_tokens=True,
        ).strip()

        print(f"Reply: {reply!r}")
        return reply or "Tell me more!"

    except Exception as e:
        import traceback
        err_str = traceback.format_exc()
        print(f"ERROR: {err_str}")
        return f"Error: {type(e).__name__}: {str(e)}"


demo = gr.ChatInterface(
    fn=respond,
    title="🤖 MiniBot — Mini Chatbot",
    description="Powered by microsoft/DialoGPT-medium",
)


if __name__ == "__main__":
    demo.launch()