Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py - HuggingFace Space for Email Classification
 import gradio as gr
 import torch
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
@@ -6,6 +5,7 @@ from setfit import SetFitModel
 import json
 import logging
 from typing import List, Dict, Any
+import os
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -20,14 +20,18 @@ def load_model():
     """Load your trained SetFit model"""
     global model, classifier
     try:
-        # Replace with your actual model path/name
+        
         model_name = "Tomiwajin/setfit_email_classifier"  
         
-        # For SetFit models
-        model = SetFitModel.from_pretrained(model_name)
-        classifier = pipeline("text-classification", model=model.model_head, tokenizer=model.model_body.tokenizer)
         
+        token = os.getenv("HF_TOKEN")
+        
+        model = SetFitModel.from_pretrained(
+            model_name, 
+            use_auth_token=token if token else True
+        )
         
+        # Create classifier directly from SetFit model
         logger.info(f"Model {model_name} loaded successfully!")
         return True
     except Exception as e:
@@ -36,22 +40,24 @@ def load_model():
 
 def classify_single_email(email_text: str) -> Dict[str, Any]:
     """Classify a single email"""
-    if not classifier:
+    if not model:
         return {"error": "Model not loaded"}
     
     try:
         # Clean and truncate text
         email_text = email_text.strip()[:5000]  # Limit length
         
-        # Get prediction
-        result = classifier(email_text)
+        # Get prediction using SetFit model directly
+        predictions = model.predict([email_text])
+        probabilities = model.predict_proba([email_text])[0]  # Get probabilities for first (and only) sample
         
-        if isinstance(result, list):
-            result = result[0]
+        # Get the predicted label and confidence
+        predicted_label = predictions[0]
+        confidence = max(probabilities)  # Confidence is the max probability
         
         return {
-            "label": result.get("label", "unknown"),
-            "score": round(result.get("score", 0.0), 4),
+            "label": str(predicted_label),
+            "score": round(float(confidence), 4),
             "success": True
         }
     except Exception as e:
@@ -60,15 +66,29 @@ def classify_single_email(email_text: str) -> Dict[str, Any]:
 
 def classify_batch_emails(emails: List[str]) -> List[Dict[str, Any]]:
     """Classify multiple emails"""
-    if not classifier:
+    if not model:
         return [{"error": "Model not loaded"}] * len(emails)
     
-    results = []
-    for email_text in emails:
-        result = classify_single_email(email_text)
-        results.append(result)
-    
-    return results
+    try:
+        # Clean and truncate texts
+        cleaned_emails = [email.strip()[:5000] for email in emails]
+        
+        # Get batch predictions
+        predictions = model.predict(cleaned_emails)
+        probabilities = model.predict_proba(cleaned_emails)
+        
+        results = []
+        for i, (pred, probs) in enumerate(zip(predictions, probabilities)):
+            results.append({
+                "label": str(pred),
+                "score": round(float(max(probs)), 4),
+                "success": True
+            })
+        
+        return results
+    except Exception as e:
+        logger.error(f"Batch classification error: {e}")
+        return [{"error": str(e), "success": False}] * len(emails)
 
 def gradio_classify(email_text: str) -> str:
     """Gradio interface function"""
@@ -97,8 +117,11 @@ def api_classify_batch(emails_json: str) -> str:
         if not isinstance(emails, list):
             return json.dumps({"error": "Input must be a JSON array of strings"})
         
+        if len(emails) > 100:  # Limit batch size
+            return json.dumps({"error": "Maximum 100 emails per batch"})
+        
         results = classify_batch_emails(emails)
-        return json.dumps(results, indent=2)
+        return json.dumps({"results": results}, indent=2)
     except json.JSONDecodeError:
         return json.dumps({"error": "Invalid JSON format"})
     except Exception as e:
@@ -151,15 +174,9 @@ with gr.Blocks(title="Email Classifier", theme=gr.themes.Soft()) as demo:
         ```
         
         ### Batch Email Classification  
-        **POST** `/api/classify-batch`
+        **POST** `/api/classify_batch`
         ```json
-        {
-            "emails": [
-                "Email 1 content...",
-                "Email 2 content...",
-                "Email 3 content..."
-            ]
-        }
+        ["Email 1 content...", "Email 2 content...", "Email 3 content..."]
         ```
         
         ### Example Response
@@ -220,48 +237,17 @@ with gr.Blocks(title="Email Classifier", theme=gr.themes.Soft()) as demo:
         const result = await response.json();
         
         // Batch classification
-        const batchResponse = await fetch('https://your-space-name.hf.space/api/classify-batch', {{
+        const batchResponse = await fetch('https://your-space-name.hf.space/api/classify_batch', {{
             method: 'POST', 
             headers: {{ 'Content-Type': 'application/json' }},
-            body: JSON.stringify({{ emails: emailArray }})
+            body: JSON.stringify(emailArray)
         }});
         const batchResults = await batchResponse.json();
         ```
         """)
 
-# Set up API endpoints
-def setup_api_routes(app):
-    """Setup FastAPI routes for the Gradio app"""
-    from fastapi import FastAPI, HTTPException
-    from pydantic import BaseModel
-    
-    class EmailRequest(BaseModel):
-        email_text: str
-    
-    class BatchEmailRequest(BaseModel):
-        emails: List[str]
-    
-    @app.post("/api/classify")
-    async def classify_endpoint(request: EmailRequest):
-        result = classify_single_email(request.email_text)
-        if not result.get("success", True):
-            raise HTTPException(status_code=500, detail=result.get("error", "Classification failed"))
-        return result
-    
-    @app.post("/api/classify-batch")
-    async def classify_batch_endpoint(request: BatchEmailRequest):
-        if len(request.emails) > 100:  # Limit batch size
-            raise HTTPException(status_code=400, detail="Maximum 100 emails per batch")
-        
-        results = classify_batch_emails(request.emails)
-        return {"results": results}
-
-# Launch the app
+# Launch the app with API endpoints
 if __name__ == "__main__":
-    # Setup API routes
-    setup_api_routes(demo.fastapi_app)
-    
-    # Launch with API support
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,