File size: 1,439 Bytes
4a74ef7 c7b12e4 4a74ef7 e2af5be 414ea47 e2af5be 414ea47 e2af5be 4a74ef7 5d73f6f 4a74ef7 5d73f6f 5d62fad 4a74ef7 002e42d 5d73f6f e2af5be a6cab09 4a74ef7 a6cab09 5d73f6f 286c3ff |
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 |
from flask import Flask, request, jsonify
import torch
import os
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
app = Flask(__name__)
# Local model path
MODEL_PATH = "./model"
print("Loading tokenizer and model from local folder:", MODEL_PATH)
# Load tokenizer & model from local folder
tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_PATH)
model = DistilBertForSequenceClassification.from_pretrained(MODEL_PATH)
model.eval()
@app.route("/predict", methods=["POST"])
def predict():
try:
data = request.get_json()
results = []
for item in data:
# Build input text
input_text = f"{item['category']} - {item['subcategory']} in {item['area']}. {item.get('comments', '')}"
inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
predicted_class = torch.argmax(outputs.logits, dim=1).item()
# Directly use predicted_class (0 = no priority)
results.append({"priority_score": predicted_class})
return jsonify(results)
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == "__main__":
# Bind to all interfaces and use port from environment (if any)
port = int(os.environ.get("PORT", 5000))
app.run(host="0.0.0.0", port=port)
|