Update app.py

app.py

@@ -1,11 +1,15 @@
 import gradio as gr
 import torch
 import re
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 
-# Load fine-tuned NER model from Hugging Face Hub
-model_name = "luxury-fashion-ner"
-ner_pipeline = pipeline("ner", model=model_name, tokenizer=model_name)
+# Load NER-Luxury model from Hugging Face
+model_name = "AkimfromParis/NER-Luxury"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForTokenClassification.from_pretrained(model_name)
+
+# Load pipeline for Named Entity Recognition (NER)
+ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
 # Regex for extracting price
 price_pattern = re.compile(r'(\bunder\b|\babove\b|\bbelow\b|\bbetween\b)?\s?(\d{1,5})\s?(AED|USD|EUR)?', re.IGNORECASE)
@@ -15,7 +19,7 @@ gender_keywords = ["men", "male", "women", "female", "unisex"]
 
 def extract_attributes(query):
     """
-    Extract structured fashion attributes dynamically using the fine-tuned NER model.
+    Extract structured fashion attributes dynamically using the fine-tuned NER-Luxury model.
     """
     structured_output = {"Brand": "Unknown", "Category": "Unknown", "Gender": "Unknown", "Price": "Unknown"}
 
@@ -26,12 +30,12 @@ def extract_attributes(query):
         entity_text = entity["word"].replace("##", "")  # Fix tokenization artifacts
         entity_label = entity["entity"]
 
-        if "ORG" in entity_label:  # Organization = Brand
+        if "HOUSE" in entity_label or "BRAND" in entity_label:  # Luxury brands
             structured_output["Brand"] = entity_text
-        elif "MISC" in entity_label:  # Miscellaneous = Category
+        elif "CATEGORY" in entity_label:  # Fashion categories
             structured_output["Category"] = entity_text
-        elif "LOC" in entity_label:  # Locations (sometimes used for brands)
-            structured_output["Brand"] = entity_text
+        elif "MONETARYVALUE" in entity_label:  # Price values
+            structured_output["Price"] = entity_text
 
     # Extract gender
     for gender in gender_keywords:
@@ -39,9 +43,9 @@ def extract_attributes(query):
             structured_output["Gender"] = gender.capitalize()
             break
 
-    # Extract price
+    # Extract price if not found by NER
     price_match = price_pattern.search(query)
-    if price_match:
+    if price_match and structured_output["Price"] == "Unknown":
         condition, amount, currency = price_match.groups()
         structured_output["Price"] = f"{condition.capitalize() if condition else ''} {amount} {currency if currency else 'AED'}".strip()
 
@@ -57,7 +61,7 @@ def parse_query(user_query):
 
 # Create Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# 🛍️ Luxury Fashion Query Parser using Fine-Tuned NER Model")
+    gr.Markdown("# 🛍️ Luxury Fashion Query Parser using NER-Luxury")
 
     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")