# ============================================================
# app.py — Gyaanchand AI Assistant (Gradio Interface)
# Author: Umer Zingu (BuzyU)
# ============================================================

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

# ============================================================
# MODEL LOADING
# ============================================================

MODEL_PATH = "UmerZingu/gyaanchand-checkpoint-10750"

# Use CUDA if available (for HF Spaces or Colab GPU)
device = 0 if torch.cuda.is_available() else -1

print(f"🚀 Loading model from {MODEL_PATH} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    low_cpu_mem_usage=True
).to("cuda" if torch.cuda.is_available() else "cpu")

# Create text-generation pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=device,
    max_new_tokens=512,
    temperature=0.7,
    top_p=0.95,
    repetition_penalty=1.1
)

# ============================================================
# CHAT FUNCTION
# ============================================================

def chat_with_gyaanchand(user_input, history=[]):
    """
    Handle chat interaction with model.
    History is used to keep multi-turn conversation.
    """
    history = history or []
    conversation = ""

    for human, ai in history:
        conversation += f"Human: {human}\nAI: {ai}\n"
    conversation += f"Human: {user_input}\nAI:"

    # Generate reply
    response = pipe(conversation)[0]["generated_text"]
    # Extract only the latest answer
    reply = response.split("AI:")[-1].strip()

    history.append((user_input, reply))
    return reply, history

# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🤖 Gyaanchand - Your Personal AI Assistant
    Talk to your fine-tuned chatbot trained from `checkpoint-10750`.

    **Built by [Umer Zingu](https://github.com/BuzyU)**  
    Running on open-source models using 🤗 Transformers + Gradio
    """)

    chatbox = gr.Chatbot(label="Chat with Gyaanchand")
    user_input = gr.Textbox(placeholder="Type your message here...", label="Your Message")

    clear = gr.Button("🧹 Clear Chat")
    send = gr.Button("Send 🚀")

    state = gr.State([])

    send.click(chat_with_gyaanchand, [user_input, state], [chatbox, state])
    clear.click(lambda: ([], []), None, [chatbox, state])

# ============================================================
# LAUNCH APP
# ============================================================

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