app.py

from flask import Flask, request, render_template
from transformers import pipeline
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"
)

def extract_entities(query):
    result = {
        "Brand": None,
        "Category": None,
        "Gender": None,
        "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"
    
    # 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"
    
    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(debug=True, host="0.0.0.0", port=7860)