56 lines
1.2 KiB
Python
56 lines
1.2 KiB
Python
from flask import Flask, request, jsonify
|
|
from flask_cors import CORS
|
|
import json
|
|
|
|
import torch
|
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
|
|
|
|
|
app = Flask(__name__)
|
|
CORS(app)
|
|
|
|
|
|
|
|
|
|
model_path = "./hate_speech_model/final_model"
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
model = AutoModelForSequenceClassification.from_pretrained(model_path)
|
|
model.eval()
|
|
|
|
@app.route("/api/predict", methods=["POST"])
|
|
def predict():
|
|
try:
|
|
data = request.json
|
|
text = data.get("text", "")
|
|
|
|
|
|
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
|
|
|
|
|
|
with torch.no_grad():
|
|
outputs = model(**inputs)
|
|
predictions = torch.argmax(outputs.logits, dim=1).item()
|
|
|
|
prediction_label = "Pravdepodobne toxický" if predictions == 1 else "Neutrálny text"
|
|
|
|
|
|
|
|
|
|
response = app.response_class(
|
|
response=json.dumps({"prediction": prediction_label}, ensure_ascii=False),
|
|
status=200,
|
|
mimetype="application/json"
|
|
)
|
|
|
|
|
|
|
|
|
|
return response
|
|
except Exception as e:
|
|
return jsonify({"error": str(e)}), 500
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0", port=5000)
|