import gradio as gr
from transformers import pipeline
import os
import spaces

MODEL_NAME = os.getenv("MODEL_NAME", "timofeyk/roberta-query-router-ecommerce")
try:
    router_pipeline = pipeline(
        "text-classification",
        model=MODEL_NAME,
        return_all_scores=True,
        device='cuda'
    )
except Exception as e:
    print(f"Error loading model: {e}")
    router_pipeline = None

@spaces.GPU
def classify_query(query_text):
    if not router_pipeline:
        raise gr.Error("Model could not be loaded. Check Space logs for details.")

    if not query_text or not query_text.strip():
        return {"Vector Search": 0.0, "Lexical Search": 0.0}

    predictions = router_pipeline(query_text)[0]

    scores = {item['label']: item['score'] for item in predictions}

    output_scores = {
        "Vector Search (Conceptual)": scores.get('vector_search', 0.0),
        "Lexical Search (Specific)": scores.get('lexical_search', 0.0)
    }

    return output_scores

title = "E-commerce Query Router"
description = """
### Is the query conceptual or specific?
Enter an e-commerce query to determine if it's better for **vector search** (conceptual, broad) or **lexical search** (specific, keyword-based). The model will output the weights for each search type.
- **Conceptual Query Example:** "bunnies themed birthday party"
- **Specific Query Example:** "all cats are beautiful t shirt"
"""
examples = [
    ["father day gift"],
    ["16x16 pillow cover"],
    ["office chair"],
    ["prom dress"],
    ["sony wh1000xm4"],
]

app = gr.Interface(
    fn=classify_query,
    inputs=gr.Textbox(
        lines=1,
        label="E-commerce Search Query",
        placeholder="Enter your product query here..."
    ),
    outputs=gr.Label(
        label="Search Type Weights",
        num_top_classes=2
    ),
    title=title,
    description=description,
    examples=examples,
    theme=gr.themes.Soft(),
    allow_flagging="never"
)

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