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

# ----------------------
# Model setup
# ----------------------
MODEL_NAME = "Dhansh2001/my-fitlien-chatbot-pruned-quantized"

def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
    
    # Ensure pad token is set
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = tokenizer.eos_token_id
    
    return tokenizer, model

tokenizer, model = load_model()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# ----------------------
# Chat function
# ----------------------
def chat_with_bot(message, history):
    try:
        # Convert history to conversation string
        conversation = ""
        for user_msg, bot_msg in history:
            conversation += f"User: {user_msg}\nBot: {bot_msg}\n"
        conversation += f"User: {message}\nBot:"
        
        # Encode input
        inputs = tokenizer.encode(conversation, return_tensors="pt").to(device)
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=500,
                num_beams=5,
                no_repeat_ngram_size=3,
                do_sample=True,
                temperature=0.7,
                pad_token_id=tokenizer.eos_token_id,
                early_stopping=True
            )
        
        # Decode only the new part
        reply = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
        return reply.strip() if reply.strip() else "I'm not sure how to answer that."
        
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# ----------------------
# Use ChatInterface instead of Blocks
# ----------------------
demo = gr.ChatInterface(
    fn=chat_with_bot,
    title="My Fitlien Chatbot",
    description="A chatbot fine-tuned from DialoGPT-medium for fitness conversations.",
    examples=[
        "Hello! How are you?",
        "What's a good workout routine?",
        "Tell me about nutrition",
        "How can I stay motivated?"
    ],
    retry_btn=None,
    undo_btn="Delete Previous",
    clear_btn="Clear",
    submit_btn="Submit"
)

# Launch
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )