# =========================================
# FLAN-T5 Chatbot (100% Stable - FINAL)
# =========================================

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

MODEL_NAME = "google/flan-t5-base"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

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


# -----------------------------
# Chat Function (IMPORTANT)
# -----------------------------
def chat(message, history):
    
    prompt = f"""
You are a helpful AI assistant.
Answer clearly and naturally.

User: {message}
Assistant:
"""

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=512
    ).to(device)

    outputs = model.generate(
        inputs.input_ids,
        max_length=120,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.2
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return response


# -----------------------------
# Gradio Chat Interface (🔥 FIX)
# -----------------------------
demo = gr.ChatInterface(
    fn=chat,
    title="🤖 AI Dialogue System (FLAN-T5)",
    description="Chat with AI using FLAN-T5"
)

demo.launch()