Update app.py

app.py

@@ -1,42 +1,20 @@
 from flask import Flask, request, jsonify
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-from transformers import BertTokenizer, BertForSequenceClassification
 import os
 
 app = Flask(__name__)
 
-# Global variables to store model and tokenizer
-global_tokenizer = None
-global_model = None
-
-def load_model():
-    """Load the model and tokenizer"""
-    global global_tokenizer, global_model
-    try:
-        print("Loading model and tokenizer...")
-        # Use a different model (bert-base-uncased)
-        MODEL_NAME = "bert-base-uncased"  # Switch to this model
-        
-        # Load tokenizer and model from Hugging Face Hub or a local path
-        global_tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
-        global_model = BertForSequenceClassification.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
-model_loaded = load_model()
+# Load the model and tokenizer directly
+tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased-finetuned-sst-2-english")
+model = AutoModelForSequenceClassification.from_pretrained("distilbert/distilbert-base-uncased-finetuned-sst-2-english")
+model.eval()  # Set the model to evaluation mode
 
 @app.route('/', methods=['GET'])
 def home():
     """Home endpoint to check if API is running"""
     response = {
         'status': 'API is running',
-        'model_status': 'loaded' if model_loaded else 'not loaded',
         'usage': {
             'endpoint': '/classify',
             'method': 'POST',
@@ -48,16 +26,11 @@ def home():
 @app.route('/health', methods=['GET'])
 def health_check():
     """Health check endpoint"""
-    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 not model_loaded:
-        return jsonify({'error': 'Model not loaded'}), 503
-
     try:
         # Get request data
         data = request.get_json()
@@ -71,11 +44,11 @@ def classify_email():
         subject = data['subject']
         
         # Tokenize
-        inputs = global_tokenizer(subject, return_tensors="pt", truncation=True, max_length=512)
+        inputs = tokenizer(subject, return_tensors="pt", truncation=True, max_length=512)
         
         # Predict
         with torch.no_grad():
-            outputs = global_model(**inputs)
+            outputs = model(**inputs)
             logits = outputs.logits
         
         # Get probabilities
@@ -85,8 +58,8 @@ def classify_email():
         
         # Define custom categories (Modify this as needed)
         CUSTOM_LABELS = {
-            0: "Business/Professional",
-            1: "Personal/Casual"
+            0: "Negative",
+            1: "Positive"
         }
         
         result = {