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

# Load base model
base_model_id = "unsloth/Qwen2.5-3B-Instruct"
adapter_model_id = "Sourabh2/qwen-fashion-assistant"

tokenizer = AutoTokenizer.from_pretrained(base_model_id)
model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    torch_dtype=torch.float16,
    device_map="auto"
)

# Load LoRA adapters
model = PeftModel.from_pretrained(model, adapter_model_id)

def chat(message, history):
    messages = [
        {
            "role": "system",
            "content": "You are a professional fashion shop consultant. Provide helpful, friendly, and knowledgeable advice about fashion, clothing, styling, and shopping."
        }
    ]

    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})

    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages,
        return_tensors="pt",
        add_generation_prompt=True
    ).to(model.device)

    outputs = model.generate(
        inputs,
        max_new_tokens=1024,
        temperature=0.6,
        top_p=0.85,
        repetition_penalty=1.1,
        do_sample=True
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = response.split("assistant")[-1].strip()

    return response

css = """
#chatbot {
    height: 600px;
}
.message {
    font-size: 16px !important;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 👔 AI Fashion Assistant
        ### Your Personal Style Consultant powered by Qwen2.5-3B
        Ask me anything about fashion, styling, outfits, colors, trends, and shopping advice!
        """
    )

    chatbot = gr.Chatbot(
        elem_id="chatbot",
        bubble_full_width=False
    )

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Ask me about fashion... (e.g., 'What should I wear to a wedding?')",
            show_label=False,
            scale=9
        )
        submit = gr.Button("Send", scale=1, variant="primary")

    gr.Examples(
        examples=[
            "What color shirt goes well with navy blue pants?",
            "I have a job interview tomorrow. What should I wear?",
            "How do I style a black leather jacket?",
            "What are the fashion trends for summer 2025?",
            "Can I wear brown shoes with a grey suit?"
        ],
        inputs=msg
    )

    msg.submit(chat, [msg, chatbot], chatbot)
    submit.click(chat, [msg, chatbot], chatbot)
    msg.submit(lambda: "", None, msg)
    submit.click(lambda: "", None, msg)

demo.launch()