Update app.py

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, jsonify, make_response
+from flask import Flask, request, jsonify
 import torch
 from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 import os
@@ -14,11 +14,8 @@ 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
-        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"
+        # Replace this path with your model's directory if using a custom model
+        MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"  # Example model, replace with your own if needed
         
         # Load tokenizer and model from Hugging Face Hub or a local path
         global_tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
@@ -86,7 +83,7 @@ def classify_email():
         predicted_class_id = logits.argmax().item()
         confidence = probabilities[0][predicted_class_id].item()
         
-        # Map to custom labels
+        # Define custom categories (Modify this as needed)
         CUSTOM_LABELS = {
             0: "Business/Professional",
             1: "Personal/Casual"