File size: 1,399 Bytes
13ec7ec f72b785 44fe967 4d95560 3596833 4d95560 44fe967 4d95560 3596833 4d95560 3596833 4d95560 3596833 4d95560 3596833 4d95560 3596833 4d95560 3596833 4d95560 16ca667 4d95560 |
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 |
import os
from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
os.environ["TRANSFORMERS_CACHE"] = "/tmp"
app = Flask(__name__)
MODEL_NAME = "s-nlp/roberta-base-formality-ranker"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
@app.route("/", methods=["GET"])
def home():
return jsonify({"message": "Formality Classifier API is running! Use /predict to classify text."})
@app.route("/predict", methods=["POST"])
def predict_formality():
data = request.get_json()
if not data or "text" not in data:
return jsonify({"error": "Text input is required"}), 400
text = data["text"]
# Tokenize input
encoding = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
# Predict formality score
with torch.no_grad():
logits = model(**encoding).logits
score = logits.softmax(dim=1)[:, 1].item()
# Convert score to formality classification
formal_percent = round(score * 100)
informal_percent = 100 - formal_percent
classification = f"Your speech is {formal_percent}% formal and {informal_percent}% informal."
return jsonify({
"text": text,
"classification": classification
})
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)
|