Add HuggingFace token

app.py

@@ -2,7 +2,7 @@ from flask import Flask, request, jsonify, render_template
 from flask_wtf import FlaskForm
 from wtforms import StringField, SubmitField
 from wtforms.validators import DataRequired
-from transformers import pipeline
+import requests
 
 app = Flask(__name__)
 app.config["SECRET_KEY"] = "secret123"
@@ -11,10 +11,12 @@ class CommentForm(FlaskForm):
     comment = StringField("Comment", validators=[DataRequired()])
     submit = SubmitField("Predict")
 
-# Load pre-trained hate speech model from Hugging Face
-print("Loading model, please wait...")
-classifier = pipeline("text-classification", model="Hate-speech-CNERG/dehatebert-mono-english")
-print("Model loaded!")
+API_URL = "https://api-inference.huggingface.co/models/Hate-speech-CNERG/dehatebert-mono-english"
+HEADERS = {"Authorization": "Bearer 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"}
+
+def query(text):
+    response = requests.post(API_URL, headers=HEADERS, json={"inputs": text})
+    return response.json()
 
 @app.route("/", methods=["GET", "POST"])
 def home():
@@ -24,17 +26,20 @@ def home():
 
     if form.validate_on_submit():
         text = form.comment.data
-        result = classifier(text)[0]
-        label = result["label"]
-        score = result["score"]
-
-        if label == "HATE":
-            prediction = "Hate Speech"
+        result = query(text)
+        if isinstance(result, list):
+            result = result[0][0]
+            label = result["label"]
+            score = result["score"]
+            if label == "HATE":
+                prediction = "Hate Speech"
+            else:
+                prediction = "No Hate Speech"
+                score = 1 - score
+            prediction_prob = round(score * 100, 2)
         else:
-            prediction = "No Hate Speech"
-            score = 1 - score
-
-        prediction_prob = round(score * 100, 2)
+            prediction = "Error - try again"
+            prediction_prob = 0
 
     return render_template("index.html", form=form, prediction=prediction, prediction_prob=prediction_prob)
 
@@ -42,33 +47,35 @@ def home():
 def predict():
     data = request.get_json()
     text = data["comment"]
-    result = classifier(text)[0]
-    label = result["label"]
-    score = result["score"]
-
-    if label == "HATE":
-        prediction = "Hate Speech"
-    else:
-        prediction = "No Hate Speech"
-        score = 1 - score
-
-    return jsonify({"prediction": prediction, "probability": round(score, 4)})
+    result = query(text)
+    if isinstance(result, list):
+        result = result[0][0]
+        label = result["label"]
+        score = result["score"]
+        if label == "HATE":
+            prediction = "Hate Speech"
+        else:
+            prediction = "No Hate Speech"
+            score = 1 - score
+        return jsonify({"prediction": prediction, "probability": round(score, 4)})
+    return jsonify({"error": "model loading"})
 
 @app.route("/prediction_by_api", methods=["GET", "POST"])
 def prediction_by_api():
     data = request.get_json()
     text = data["comment"]
-    result = classifier(text)[0]
-    label = result["label"]
-    score = result["score"]
-
-    if label == "HATE":
-        prediction = "Hate Speech"
-    else:
-        prediction = "No Hate Speech"
-        score = 1 - score
-
-    return jsonify({"prediction": prediction, "probability": round(score, 4)})
+    result = query(text)
+    if isinstance(result, list):
+        result = result[0][0]
+        label = result["label"]
+        score = result["score"]
+        if label == "HATE":
+            prediction = "Hate Speech"
+        else:
+            prediction = "No Hate Speech"
+            score = 1 - score
+        return jsonify({"prediction": prediction, "probability": round(score, 4)})
+    return jsonify({"error": "model loading"})
 
 if __name__ == "__main__":
     app.run(debug=True)