Init

app.py

@@ -0,0 +1,71 @@
+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
+    )
+    router_pipeline.to('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:
+        return {"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:** "summer vibes clothing"
+- **Specific Query Example:** "nike air force 1 size 10"
+"""
+examples = [
+    ["father day gift"],
+    ["16x16 pillow cover"],
+    ["something to wear for a wedding"],
+    ["logitech mx master 3s mouse"],
+    ["comfortable office chair"],
+]
+
+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()