Update app.py

app.py

@@ -1,76 +1,35 @@
+# app.py
 from flask import Flask, request, render_template
-from huggingface_hub import InferenceClient
+from transformers import pipeline
 import re
 
 app = Flask(__name__)
 
-# Initialize DeepSeek-R1 client
-client = InferenceClient(model="deepseek-ai/deepseek-llm-67b-chat")
+# Load lightweight fashion model
+ner_pipeline = pipeline(
+    "ner",
+    model="Keng/bert-product-ner",
+    aggregation_strategy="simple"
+)
 
-def parse_llm_response(response):
-    """Improved parsing that handles model's raw responses"""
-    result = {
-        "Brand": None,
-        "Category": None,
-        "Gender": None,
-        "Price": None
-    }
+def extract_details(query):
+    result = {"Brand": None, "Category": None, "Price": None}
     
-    # Enhanced pattern matching for flexible JSON extraction
-    patterns = {
-        "brand": r'"brand":\s*"([^"]*)"',
-        "category": r'"category":\s*"([^"]*)"',
-        "gender": r'"gender":\s*"([^"]*)"',
-        "price_range": r'"price_range":\s*"([^"]*)"'
-    }
+    # 1. NER Extraction
+    entities = ner_pipeline(query)
+    for entity in entities:
+        if entity["entity_group"] == "BRAND":
+            result["Brand"] = entity["word"].title()
+        elif entity["entity_group"] == "PRODUCT":
+            result["Category"] = entity["word"].title()
     
-    for key, pattern in patterns.items():
-        match = re.search(pattern, response, re.IGNORECASE)
-        if match:
-            value = match.group(1).strip()
-            if value.lower() in ["null", "n/a", ""]:
-                continue
-            if key == "brand":
-                result["Brand"] = value.title()
-            elif key == "category":
-                result["Category"] = value.title()
-            elif key == "gender":
-                result["Gender"] = value.title()
-            elif key == "price_range":
-                result["Price"] = value.upper()
+    # 2. Price Regex Fallback
+    price_match = re.search(r"(under|below|less than)\s*(\d+)\s*AED", query, re.IGNORECASE)
+    if price_match:
+        result["Price"] = f"{price_match.group(1).title()} {price_match.group(2)} AED"
     
     return result
 
-def analyze_query(query):
-    """Enhanced prompt for luxury brand understanding"""
-    prompt = f"""Analyze this fashion query and extract structured data. Follow these rules:
-
-1. Brand: Identify the luxury fashion brand mentioned (e.g., Gucci, Prada, Balenciaga)
-2. Category: Product type (perfume, bag, shoes, etc.)
-3. Gender: men, women, or unisex
-4. Price: Exact price range from query
-
-Return JSON format:
-
-{{
-    "brand": "<brand name>",
-    "category": "<product category>",
-    "gender": "<target gender>",
-    "price_range": "<price info>"
-}}
-
-Query: "{query}"
-"""
-
-    response = client.text_generation(
-        prompt=prompt,
-        max_new_tokens=200,
-        temperature=0.01,  # More deterministic output
-        stop_sequences=["\n\n"]  # Prevent extra text
-    )
-    
-    return parse_llm_response(response)
-
 @app.route("/", methods=["GET", "POST"])
 def index():
     result = None
@@ -78,7 +37,7 @@ def index():
     if request.method == "POST":
         query = request.form.get("query", "")
         if query.strip():
-            result = analyze_query(query)
+            result = extract_details(query)
     return render_template("index.html", result=result, query=query)
 
 if __name__ == "__main__":