Davephoenix commited on
Commit
3c061e1
·
1 Parent(s): 7afa047

Small change

Browse files
Files changed (1) hide show
  1. app.py +7 -5
app.py CHANGED
@@ -3,16 +3,15 @@ from transformers import pipeline
3
 
4
  app = Flask(__name__)
5
 
6
- # Define label mappings
7
- id2label = {0: "not_bullying", 1: "bullying"}
8
- label2id = {"not_bullying": 0, "bullying": 1}
 
9
 
10
  # Load the text classification pipeline with label mappings
11
  classifier = pipeline(
12
  "text-classification",
13
  model="Davephoenix/bert-bullying-detector",
14
- id2label=id2label,
15
- label2id=label2id,
16
  )
17
 
18
 
@@ -23,6 +22,9 @@ def classify_text():
23
  return jsonify({"error": "Missing 'text' in request body"}), 400
24
  try:
25
  result = classifier(data["text"])
 
 
 
26
  return jsonify(result)
27
  except Exception as e:
28
  return jsonify({"error": str(e)}), 500
 
3
 
4
  app = Flask(__name__)
5
 
6
+ labels = {
7
+ "LABEL_0": "not_bullying",
8
+ "LABEL_1": "bullying",
9
+ }
10
 
11
  # Load the text classification pipeline with label mappings
12
  classifier = pipeline(
13
  "text-classification",
14
  model="Davephoenix/bert-bullying-detector",
 
 
15
  )
16
 
17
 
 
22
  return jsonify({"error": "Missing 'text' in request body"}), 400
23
  try:
24
  result = classifier(data["text"])
25
+ # Map the labels to their meanings
26
+ for item in result:
27
+ item["label"] = labels.get(item["label"], item["label"])
28
  return jsonify(result)
29
  except Exception as e:
30
  return jsonify({"error": str(e)}), 500