QueryAnalyzerV2 / app.py
DINGOLANI's picture
Update app.py
57b32e0 verified
raw
history blame
2.45 kB
import gradio as gr
import torch
import re
from transformers import 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)
# 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)
# Keywords for gender extraction
gender_keywords = ["men", "male", "women", "female", "unisex"]
def extract_attributes(query):
"""
Extract structured fashion attributes dynamically using the fine-tuned NER model.
"""
structured_output = {"Brand": "Unknown", "Category": "Unknown", "Gender": "Unknown", "Price": "Unknown"}
# Run NER model on query
entities = ner_pipeline(query)
for entity in entities:
entity_text = entity["word"].replace("##", "") # Fix tokenization artifacts
entity_label = entity["entity"]
if "ORG" in entity_label: # Organization = Brand
structured_output["Brand"] = entity_text
elif "MISC" in entity_label: # Miscellaneous = Category
structured_output["Category"] = entity_text
elif "LOC" in entity_label: # Locations (sometimes used for brands)
structured_output["Brand"] = entity_text
# Extract gender
for gender in gender_keywords:
if gender in query.lower():
structured_output["Gender"] = gender.capitalize()
break
# Extract price
price_match = price_pattern.search(query)
if price_match:
condition, amount, currency = price_match.groups()
structured_output["Price"] = f"{condition.capitalize() if condition else ''} {amount} {currency if currency else 'AED'}".strip()
return structured_output
# Define Gradio UI
def parse_query(user_query):
"""
Parses fashion-related queries into structured attributes.
"""
parsed_output = extract_attributes(user_query)
return parsed_output # JSON output
# Create Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("# 🛍️ Luxury Fashion Query Parser using Fine-Tuned NER Model")
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")
parse_button = gr.Button("Parse Query")
parse_button.click(parse_query, inputs=[query_input], outputs=[output_box])
demo.launch()