Update app.py

app.py

@@ -14,10 +14,14 @@ def load_model():
     global global_tokenizer, global_model
     try:
         print("Loading model and tokenizer...")
-        MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"
+        # Replace this path with your model's directory if using a custom model
+        MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"  # Can be a custom path if using your own model
+        
+        # Load tokenizer and model from Hugging Face Hub or a local path
         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:
@@ -25,14 +29,14 @@ def load_model():
         return False
 
 # Load model at startup
-load_model()
+model_loaded = 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',
+        'model_status': 'loaded' if model_loaded else 'not loaded',
         'usage': {
             'endpoint': '/classify',
             'method': 'POST',
@@ -44,14 +48,14 @@ def home():
 @app.route('/health', methods=['GET'])
 def health_check():
     """Health check endpoint"""
-    if global_model is None or global_tokenizer is None:
+    if not model_loaded:
         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:
+    if not model_loaded:
         return jsonify({'error': 'Model not loaded'}), 503
 
     try:
@@ -101,6 +105,6 @@ def classify_email():
         return jsonify({'error': str(e)}), 500
 
 if __name__ == '__main__':
-    # Use port 7860 for Hugging Face Spaces
+    # Use port 7860 for Hugging Face Spaces or any other port for local testing
     port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='0.0.0.0', port=port)