Update app.py

app.py

@@ -1,67 +1,21 @@
-from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 from flask import Flask, request, jsonify
-import torch
+from model import load_model, predict_deal_score
+import logging
 
-# Initialize Flask app
 app = Flask(__name__)
 
-# Load pre-trained DistilBERT model and tokenizer
-model_name = "distilbert-base-uncased"
-tokenizer = DistilBertTokenizer.from_pretrained(model_name)
-model = DistilBertForSequenceClassification.from_pretrained(model_name)
-
-# Ensure the model is in evaluation mode
-model.eval()
-
-# Define a function to predict score and risk based on the input
-def predict_deal_qualification(inputs):
-    # Prepare the input text (combining fields for the model)
-    input_text = f"Industry: {inputs['industry']}, Stage: {inputs['stage']}, Amount: {inputs['amount']}, Lead Score: {inputs['lead_score']}, Email count: {inputs['emails_last_7_days']}, Meeting count: {inputs['meetings_last_30_days']}"
-
-    # Tokenize the input text
-    tokens = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True, max_length=128)
-
-    # Run model inference
-    with torch.no_grad():
-        outputs = model(**tokens)
-
-    # Get prediction logits
-    logits = outputs.logits
-
-    # Convert logits to probability (0-1 range)
-    score = torch.sigmoid(logits).item()  # Calculate the probability
-    score = round(score * 100, 2)  # Scale the score to 0-100
-
-    # Map logits to risk classification (Low, Medium, High)
-    risk_classes = ["Low", "Medium", "High"]
-    risk = risk_classes[torch.argmax(logits).item()]
-
-    # Generate a dummy recommendation (you can customize this based on business rules)
-    recommendation = "Schedule another meeting before sending proposal."  # Replace with your custom recommendation logic
-
-    return {
-        "score": score,
-        "confidence": round(torch.max(torch.softmax(logits, dim=-1)).item(), 2),
-        "risk": risk,
-        "recommendation": recommendation
-    }
-
-# Define an endpoint for deal qualification prediction
-@app.route('/predict', methods=['POST'])
-def predict():
-    # Get input JSON data from the POST request
-    data = request.get_json()
-
-    # Validate input structure (ensure required fields are present)
-    if not all(key in data for key in ['industry', 'stage', 'amount', 'lead_score', 'emails_last_7_days', 'meetings_last_30_days']):
-        return jsonify({"error": "Missing required input data"}), 400
-
-    # Predict using the pre-trained model
-    result = predict_deal_qualification(data)
-
-    # Return the prediction result as a JSON response
-    return jsonify(result)
-
-# Run the Flask app
-if __name__ == '__main__':
-    app.run(debug=True, host="0.0.0.0", port=5000)
+# Load model once at startup
+model, tokenizer, summarizer = load_model()
+
+@app.route("/score", methods=["POST"])
+def score():
+    try:
+        data = request.get_json()
+        result = predict_deal_score(data, model, tokenizer, summarizer)
+        return jsonify(result), 200
+    except Exception as e:
+        logging.exception("Prediction error")
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)