app.py

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__":
    import gunicorn.app.base

    class StandaloneApplication(gunicorn.app.base.BaseApplication):
        def __init__(self, app, options=None):
            self.options = options or {}
            self.application = app
            super().__init__()

        def load_config(self):
            for key, value in self.options.items():
                self.cfg.set(key.lower(), value)

        def load(self):
            return self.application

    options = {
        "bind": "0.0.0.0:7860",
        "workers": 1,  # Increase if you have more CPU/GPU
        "threads": 2,
        "timeout": 120,  # seconds to allow startup
    }
    StandaloneApplication(app, options).run()