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

# -------------------------------------------------
# Configuration
# -------------------------------------------------
BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
LORA_REPO = "Srikanthgoud7/qwen2.5-customer-support-lora"  

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# -------------------------------------------------
# Load tokenizer & model
# -------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

model = PeftModel.from_pretrained(base_model, LORA_REPO)
model.eval()

# -------------------------------------------------
# Chat function (DYNAMIC INPUT)
# -------------------------------------------------
def chat(user_input, history):
    """
    user_input: current user message
    history: list of (user, assistant) tuples
    """

    # Build conversation prompt
    conversation = ""
    for u, a in history:
        conversation += f"User: {u}\nAssistant: {a}\n"
    conversation += f"User: {user_input}\nAssistant:"

    inputs = tokenizer(conversation, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True
        )

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

    # Extract only the assistant's latest reply
    response = response.split("Assistant:")[-1].strip()

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

# -------------------------------------------------
# Gradio Chat UI
# -------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 Customer Support Chatbot (LoRA + Qwen2.5)")

    chatbot = gr.Chatbot()
    user_input = gr.Textbox(
        placeholder="Type your query here (e.g., Can you cancel my order?)",
        label="Your Message"
    )
    state = gr.State([])

    send_btn = gr.Button("Send")
    clear_btn = gr.Button("Clear Chat")

    send_btn.click(
        fn=chat,
        inputs=[user_input, state],
        outputs=[chatbot, state]
    )

    clear_btn.click(
        fn=lambda: ([], []),
        inputs=None,
        outputs=[chatbot, state]
    )

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