DINGOLANI commited on
Commit
88c3dd8
·
verified ·
1 Parent(s): 7cc7d48

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -0
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)