Update app.py

app.py

@@ -1,17 +1,14 @@
 from flask import Flask, request, render_template
-from transformers import pipeline
+from huggingface_hub import InferenceClient
 import re
 
 app = Flask(__name__)
 
-# Use a PUBLICLY AVAILABLE model that works on free tier
-ner_pipeline = pipeline(
-    "ner",
-    model="dslim/bert-base-NER",  # Verified public model
-    aggregation_strategy="simple"
-)
+# Initialize DeepSeek-R1 client
+client = InferenceClient(model="deepseek-ai/deepseek-llm-67b-chat")
 
-def extract_entities(query):
+def parse_llm_response(response):
+    """Improved parsing that handles model's raw responses"""
     result = {
         "Brand": None,
         "Category": None,
@@ -19,37 +16,70 @@ def extract_entities(query):
         "Price": None
     }
     
-    # Extract entities
-    entities = ner_pipeline(query)
-    
-    # Process entities
-    for entity in entities:
-        if entity["entity_group"] == "ORG":  # Organizations are likely brands
-            result["Brand"] = entity["word"]
-    
-    # Add keyword-based extraction for other fields
-    query_lower = query.lower()
-    if "perfume" in query_lower or "cologne" in query_lower:
-        result["Category"] = "Perfume"
-    if "men" in query_lower:
-        result["Gender"] = "Men"
-    elif "women" in query_lower:
-        result["Gender"] = "Women"
+    # 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*"([^"]*)"'
+    }
     
-    # Price extraction
-    price_match = re.search(r"under (\d+)\s*AED", query, re.IGNORECASE)
-    if price_match:
-        result["Price"] = f"Under {price_match.group(1)} AED"
+    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()
     
     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
+    query = ""
     if request.method == "POST":
-        query = request.form["query"]
-        result = extract_entities(query)
-        return render_template("index.html", result=result, query=query)
-    return render_template("index.html", result=None)
+        query = request.form.get("query", "")
+        if query.strip():
+            result = analyze_query(query)
+    return render_template("index.html", result=result, query=query)
 
 if __name__ == "__main__":
-    app.run(debug=True, host="0.0.0.0", port=7860)
+    app.run(host="0.0.0.0", port=7860)