import gradio as gr
import torch
from transformers import CLIPProcessor, CLIPModel

# Load the FashionCLIP model
model_name = "patrickjohncyh/fashion-clip"
model = CLIPModel.from_pretrained(model_name)
processor = CLIPProcessor.from_pretrained(model_name)

def parse_query(user_query):
    """
    Parse fashion-related search queries into structured data.
    """
    # Define categories relevant to luxury fashion search
    fashion_categories = ["Brand", "Category", "Gender", "Price Range"]

    # Format user query for CLIP
    inputs = processor(text=[user_query], images=None, return_tensors="pt", padding=True)

    # Get model embeddings
    with torch.no_grad():
        outputs = model.get_text_features(**inputs)

    # Simulated parsing output (FashionCLIP itself does not generate structured JSON)
    parsed_output = {
        "Brand": "Gucci" if "Gucci" in user_query else "Unknown",
        "Category": "Perfume" if "perfume" in user_query else "Unknown",
        "Gender": "Men" if "men" in user_query else "Women" if "women" in user_query else "Unisex",
        "Price Range": "Under 200 AED" if "under 200" in user_query else "Above 200 AED",
    }

    return parsed_output

# Define Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🛍️ Luxury Fashion Query Parser (FashionCLIP)")
    
    query_input = gr.Textbox(label="Enter your search query", placeholder="e.g., Gucci men’s perfume under 200AED")
    output_box = gr.JSON(label="Parsed Output")

    parse_button = gr.Button("Parse Query")
    parse_button.click(parse_query, inputs=[query_input], outputs=[output_box])

# Launch the app
demo.launch()