Create app.py

app.py

@@ -0,0 +1,79 @@
+from flask import Flask, request, jsonify
+import torch
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
+import os
+
+app = Flask(__name__)
+
+# Initialize model and tokenizer globally
+print("Loading model and tokenizer...")
+MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"
+tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
+model = DistilBertForSequenceClassification.from_pretrained(MODEL_NAME)
+model.eval()
+print("Model and tokenizer loaded successfully!")
+
+# Custom labels
+CUSTOM_LABELS = {
+    0: "Business/Professional",
+    1: "Personal/Casual"
+}
+
+def classify_text(text):
+    try:
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+        
+        probabilities = torch.nn.functional.softmax(logits, dim=1)
+        predicted_class_id = logits.argmax().item()
+        confidence = probabilities[0][predicted_class_id].item()
+        
+        return {
+            'category': CUSTOM_LABELS[predicted_class_id],
+            'confidence': round(confidence, 3),
+            'all_categories': {
+                label: round(prob.item(), 3) 
+                for label, prob in zip(CUSTOM_LABELS.values(), probabilities[0])
+            }
+        }
+    except Exception as e:
+        print(f"Error in classify_text: {str(e)}")
+        raise
+
+@app.route('/classify', methods=['POST'])
+def classify_email():
+    try:
+        data = request.get_json()
+        
+        if not data or 'subject' not in data:
+            return jsonify({
+                'error': 'No subject provided. Please send a JSON with "subject" field.'
+            }), 400
+        
+        subject = data['subject']
+        result = classify_text(subject)
+        return jsonify(result)
+    
+    except Exception as e:
+        print(f"Error in classify_email: {str(e)}")
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/', methods=['GET'])
+def home():
+    return jsonify({
+        'status': 'API is running',
+        'model_name': MODEL_NAME,
+        'usage': {
+            'endpoint': '/classify',
+            'method': 'POST',
+            'body': {'subject': 'Your email subject here'}
+        }
+    })
+
+if __name__ == '__main__':
+    port = int(os.environ.get('PORT', 7860))
+    print(f"Starting server on port {port}...")
+    app.run(host='0.0.0.0', port=port, debug=True)