Spaces:
Runtime error
Runtime error
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,54 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from flask import Flask, request, render_template
|
2 |
+
from transformers import pipeline
|
3 |
+
import re
|
4 |
+
|
5 |
+
app = Flask(__name__)
|
6 |
+
|
7 |
+
# Load the NER pipeline
|
8 |
+
ner_pipeline = pipeline(
|
9 |
+
"ner",
|
10 |
+
model="VoiceLab/bert-base-uncased-finetuned-ner-fashion",
|
11 |
+
aggregation_strategy="simple"
|
12 |
+
)
|
13 |
+
|
14 |
+
def extract_entities(query):
|
15 |
+
# Get entities from model
|
16 |
+
entities = ner_pipeline(query)
|
17 |
+
|
18 |
+
# Initialize result dictionary
|
19 |
+
result = {
|
20 |
+
"Brand": None,
|
21 |
+
"Category": None,
|
22 |
+
"Gender": None,
|
23 |
+
"Price": None
|
24 |
+
}
|
25 |
+
|
26 |
+
# Process entities
|
27 |
+
for entity in entities:
|
28 |
+
label = entity["entity_group"]
|
29 |
+
value = entity["word"].strip()
|
30 |
+
|
31 |
+
if label == "Brand" and not result["Brand"]:
|
32 |
+
result["Brand"] = value
|
33 |
+
elif label == "Product Type" and not result["Category"]:
|
34 |
+
result["Category"] = value
|
35 |
+
elif label == "Gender" and not result["Gender"]:
|
36 |
+
result["Gender"] = value
|
37 |
+
|
38 |
+
# Extract price using regex (since model doesn't handle prices)
|
39 |
+
price_match = re.search(r"under (\d+)\s*AED", query, re.IGNORECASE)
|
40 |
+
if price_match:
|
41 |
+
result["Price"] = f"Under {price_match.group(1)} AED"
|
42 |
+
|
43 |
+
return result
|
44 |
+
|
45 |
+
@app.route("/", methods=["GET", "POST"])
|
46 |
+
def index():
|
47 |
+
if request.method == "POST":
|
48 |
+
query = request.form["query"]
|
49 |
+
result = extract_entities(query)
|
50 |
+
return render_template("index.html", result=result, query=query)
|
51 |
+
return render_template("index.html", result=None)
|
52 |
+
|
53 |
+
if __name__ == "__main__":
|
54 |
+
app.run(host="0.0.0.0", port=5000, debug=True)
|