Update app.py

app.py

@@ -14,8 +14,8 @@ def load_model():
     global global_tokenizer, global_model
     try:
         print("Loading model and tokenizer...")
-        # 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
+        # Replace this path with your model's directory or Hugging Face model name
+        MODEL_NAME = "aideveloper24/email_classify"  # Replace with your custom model name
         
         # Load tokenizer and model from Hugging Face Hub or a local path
         global_tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
@@ -70,10 +70,10 @@ def classify_email():
         # Get the subject
         subject = data['subject']
         
-        # Tokenize
+        # Tokenize the subject text
         inputs = global_tokenizer(subject, return_tensors="pt", truncation=True, max_length=512)
         
-        # Predict
+        # Predict the class
         with torch.no_grad():
             outputs = global_model(**inputs)
             logits = outputs.logits
@@ -89,6 +89,7 @@ def classify_email():
             1: "Personal/Casual"
         }
         
+        # Create the response
         result = {
             'category': CUSTOM_LABELS[predicted_class_id],
             'confidence': round(confidence, 3),