Spaces:
Sleeping
Sleeping
File size: 2,783 Bytes
1e8265a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
from flask import Flask, request, jsonify, render_template
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F
import re
from flask_cors import CORS # Enable CORS
# Initialize Flask app
app = Flask(__name__)
CORS(app) # Allow requests from frontend apps
# Choose your model: 'bert-base-uncased' or 'GroNLP/hateBERT'
MODEL_NAME = 'bert-base-uncased' # Change to 'GroNLP/hateBERT' if needed
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()
# Two-class labels only
LABELS = ['Safe', 'Cyberbullying']
# Offensive trigger words
TRIGGER_WORDS = [
"gago", "pokpok", "yawa", "linte", "ulol", "tangina", "bilat", "putang", "tarantado", "bobo",
"yudipota", "law-ay", "bilatibay", "hayop"
]
# Detect trigger words in input text
def find_triggers(text):
found = []
for word in TRIGGER_WORDS:
if re.search(rf"\b{re.escape(word)}\b", text, re.IGNORECASE):
found.append(word)
return found
# Predict function
def predict_text(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
# Use GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
inputs = {key: value.to(device) for key, value in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probs = F.softmax(logits, dim=1)
confidence, predicted_class = torch.max(probs, dim=1)
# Fallback logic: if model predicts more than 2 classes, default to Safe if out-of-bounds
label_index = predicted_class.item()
if label_index >= len(LABELS):
label_index = 0 # default to "Safe"
label = LABELS[label_index]
confidence_score = round(confidence.item(), 4)
triggers = find_triggers(text)
# Override model prediction if offensive triggers found
if triggers and label == "Safe":
label = "Cyberbullying"
return {
"label": label,
"confidence": confidence_score,
"triggers": triggers
}
# Serve frontend
@app.route('/')
def index():
return render_template('index.html') # Ensure templates/index.html exists
# API endpoint
@app.route("/predict", methods=["POST"])
def predict_api():
try:
data = request.get_json()
text = data.get("text", "")
if not text.strip():
return jsonify({"error": "No text provided"}), 400
result = predict_text(text)
return jsonify(result)
except Exception as e:
return jsonify({"error": str(e)}), 500
# Run server
#if __name__ == "__main__":
# app.run(debug=True)
|