app.py

from flask import Flask, request, render_template
from transformers import pipeline
import re

app = Flask(__name__)

# Load the NER pipeline
ner_pipeline = pipeline(
    "ner",
    model="VoiceLab/bert-base-uncased-finetuned-ner-fashion",
    aggregation_strategy="simple"
)

def extract_entities(query):
    # Get entities from model
    entities = ner_pipeline(query)
    
    # Initialize result dictionary
    result = {
        "Brand": None,
        "Category": None,
        "Gender": None,
        "Price": None
    }
    
    # Process entities
    for entity in entities:
        label = entity["entity_group"]
        value = entity["word"].strip()
        
        if label == "Brand" and not result["Brand"]:
            result["Brand"] = value
        elif label == "Product Type" and not result["Category"]:
            result["Category"] = value
        elif label == "Gender" and not result["Gender"]:
            result["Gender"] = value
    
    # Extract price using regex (since model doesn't handle prices)
    price_match = re.search(r"under (\d+)\s*AED", query, re.IGNORECASE)
    if price_match:
        result["Price"] = f"Under {price_match.group(1)} AED"
    
    return result

@app.route("/", methods=["GET", "POST"])
def index():
    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)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)