DINGOLANI commited on
Commit
5e0ca02
·
verified ·
1 Parent(s): aa81da0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -17
app.py CHANGED
@@ -4,18 +4,14 @@ 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,
@@ -23,19 +19,24 @@ def extract_entities(query):
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"
@@ -51,4 +52,4 @@ def index():
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)
 
4
 
5
  app = Flask(__name__)
6
 
7
+ # Use a PUBLICLY AVAILABLE model that works on free tier
8
  ner_pipeline = pipeline(
9
  "ner",
10
+ model="dslim/bert-base-NER", # Verified public model
11
  aggregation_strategy="simple"
12
  )
13
 
14
  def extract_entities(query):
 
 
 
 
15
  result = {
16
  "Brand": None,
17
  "Category": None,
 
19
  "Price": None
20
  }
21
 
22
+ # Extract entities
23
+ entities = ner_pipeline(query)
24
+
25
  # Process entities
26
  for entity in entities:
27
+ if entity["entity_group"] == "ORG": # Organizations are likely brands
28
+ result["Brand"] = entity["word"]
29
+
30
+ # Add keyword-based extraction for other fields
31
+ query_lower = query.lower()
32
+ if "perfume" in query_lower or "cologne" in query_lower:
33
+ result["Category"] = "Perfume"
34
+ if "men" in query_lower:
35
+ result["Gender"] = "Men"
36
+ elif "women" in query_lower:
37
+ result["Gender"] = "Women"
38
 
39
+ # Price extraction
40
  price_match = re.search(r"under (\d+)\s*AED", query, re.IGNORECASE)
41
  if price_match:
42
  result["Price"] = f"Under {price_match.group(1)} AED"
 
52
  return render_template("index.html", result=None)
53
 
54
  if __name__ == "__main__":
55
+ app.run(debug=True, host="0.0.0.0", port=7860)