from flask import Flask, request, jsonify
import torch
import os
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification

app = Flask(__name__)

MODEL_PATH = "./model"
tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_PATH + "/tokenizer", local_files_only=True)
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:
            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()

            results.append({"priority_score": predicted_class})

        return jsonify(results)
    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    app.run(port=5000)