Create app.py

app.py

@@ -1,37 +1,91 @@
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, make_response
 import torch
 from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 import os
 
 app = Flask(__name__)
 
-# Initialize model and tokenizer globally
-print("Loading model and tokenizer...")
-MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"
-tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
-model = DistilBertForSequenceClassification.from_pretrained(MODEL_NAME)
-model.eval()
-print("Model and tokenizer loaded successfully!")
+# Global variables to store model and tokenizer
+global_tokenizer = None
+global_model = None
 
-# Custom labels
-CUSTOM_LABELS = {
-    0: "Business/Professional",
-    1: "Personal/Casual"
-}
+def load_model():
+    """Load the model and tokenizer"""
+    global global_tokenizer, global_model
+    try:
+        print("Loading model and tokenizer...")
+        MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"
+        global_tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
+        global_model = DistilBertForSequenceClassification.from_pretrained(MODEL_NAME)
+        global_model.eval()
+        print("Model loaded successfully!")
+        return True
+    except Exception as e:
+        print(f"Error loading model: {str(e)}")
+        return False
+
+# Load model at startup
+load_model()
+
+@app.route('/', methods=['GET'])
+def home():
+    """Home endpoint to check if API is running"""
+    response = {
+        'status': 'API is running',
+        'model_status': 'loaded' if global_model is not None else 'not loaded',
+        'usage': {
+            'endpoint': '/classify',
+            'method': 'POST',
+            'body': {'subject': 'Your email subject here'}
+        }
+    }
+    return jsonify(response)
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    """Health check endpoint"""
+    if global_model is None or global_tokenizer is None:
+        return jsonify({'status': 'unhealthy', 'error': 'Model not loaded'}), 503
+    return jsonify({'status': 'healthy'})
+
+@app.route('/classify', methods=['POST'])
+def classify_email():
+    """Classify email subject"""
+    if global_model is None or global_tokenizer is None:
+        return jsonify({'error': 'Model not loaded'}), 503
 
-def classify_text(text):
     try:
-        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+        # Get request data
+        data = request.get_json()
+        
+        if not data or 'subject' not in data:
+            return jsonify({
+                'error': 'No subject provided. Please send a JSON with "subject" field.'
+            }), 400
+        
+        # Get the subject
+        subject = data['subject']
         
+        # Tokenize
+        inputs = global_tokenizer(subject, return_tensors="pt", truncation=True, max_length=512)
+        
+        # Predict
         with torch.no_grad():
-            outputs = model(**inputs)
+            outputs = global_model(**inputs)
             logits = outputs.logits
         
+        # Get probabilities
         probabilities = torch.nn.functional.softmax(logits, dim=1)
         predicted_class_id = logits.argmax().item()
         confidence = probabilities[0][predicted_class_id].item()
         
-        return {
+        # Map to custom labels
+        CUSTOM_LABELS = {
+            0: "Business/Professional",
+            1: "Personal/Casual"
+        }
+        
+        result = {
             'category': CUSTOM_LABELS[predicted_class_id],
             'confidence': round(confidence, 3),
             'all_categories': {
@@ -39,41 +93,14 @@ def classify_text(text):
                 for label, prob in zip(CUSTOM_LABELS.values(), probabilities[0])
             }
         }
-    except Exception as e:
-        print(f"Error in classify_text: {str(e)}")
-        raise
-
-@app.route('/classify', methods=['POST'])
-def classify_email():
-    try:
-        data = request.get_json()
-        
-        if not data or 'subject' not in data:
-            return jsonify({
-                'error': 'No subject provided. Please send a JSON with "subject" field.'
-            }), 400
         
-        subject = data['subject']
-        result = classify_text(subject)
         return jsonify(result)
     
     except Exception as e:
-        print(f"Error in classify_email: {str(e)}")
+        print(f"Error in classification: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
-@app.route('/', methods=['GET'])
-def home():
-    return jsonify({
-        'status': 'API is running',
-        'model_name': MODEL_NAME,
-        'usage': {
-            'endpoint': '/classify',
-            'method': 'POST',
-            'body': {'subject': 'Your email subject here'}
-        }
-    })
-
 if __name__ == '__main__':
+    # Use port 7860 for Hugging Face Spaces
     port = int(os.environ.get('PORT', 7860))
-    print(f"Starting server on port {port}...")
-    app.run(host='0.0.0.0', port=port, debug=True)
+    app.run(host='0.0.0.0', port=port)