Update app.py

app.py

@@ -14,8 +14,11 @@ def load_model():
     global global_tokenizer, global_model
     try:
         print("Loading model and tokenizer...")
-        # Replace this path with your model's directory or Hugging Face model name
-        MODEL_NAME = "aideveloper24/email_classify"  # Replace with your custom model name
+        # Replace this path with your model's directory or Hugging Face model
+        MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"  # Test with a known model for now
+
+        # If you have a local model path, use the path to your model
+        # model_dir = "/path/to/your/local/model"
         
         # Load tokenizer and model from Hugging Face Hub or a local path
         global_tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
@@ -70,10 +73,10 @@ def classify_email():
         # Get the subject
         subject = data['subject']
         
-        # Tokenize the subject text
+        # Tokenize
         inputs = global_tokenizer(subject, return_tensors="pt", truncation=True, max_length=512)
         
-        # Predict the class
+        # Predict
         with torch.no_grad():
             outputs = global_model(**inputs)
             logits = outputs.logits
@@ -89,7 +92,6 @@ def classify_email():
             1: "Personal/Casual"
         }
         
-        # Create the response
         result = {
             'category': CUSTOM_LABELS[predicted_class_id],
             'confidence': round(confidence, 3),