Create app.py

app.py

@@ -0,0 +1,270 @@
+# app.py - HuggingFace Space for Email Classification
+import gradio as gr
+import torch
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+from setfit import SetFitModel
+import json
+import logging
+from typing import List, Dict, Any
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global model variable
+model = None
+tokenizer = None
+classifier = None
+
+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)
+        
+        
+        logger.info(f"Model {model_name} loaded successfully!")
+        return True
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        return False
+
+def classify_single_email(email_text: str) -> Dict[str, Any]:
+    """Classify a single email"""
+    if not classifier:
+        return {"error": "Model not loaded"}
+    
+    try:
+        # Clean and truncate text
+        email_text = email_text.strip()[:5000]  # Limit length
+        
+        # Get prediction
+        result = classifier(email_text)
+        
+        if isinstance(result, list):
+            result = result[0]
+        
+        return {
+            "label": result.get("label", "unknown"),
+            "score": round(result.get("score", 0.0), 4),
+            "success": True
+        }
+    except Exception as e:
+        logger.error(f"Classification error: {e}")
+        return {"error": str(e), "success": False}
+
+def classify_batch_emails(emails: List[str]) -> List[Dict[str, Any]]:
+    """Classify multiple emails"""
+    if not classifier:
+        return [{"error": "Model not loaded"}] * len(emails)
+    
+    results = []
+    for email_text in emails:
+        result = classify_single_email(email_text)
+        results.append(result)
+    
+    return results
+
+def gradio_classify(email_text: str) -> str:
+    """Gradio interface function"""
+    if not email_text.strip():
+        return "Please enter some email text to classify."
+    
+    result = classify_single_email(email_text)
+    
+    if result.get("success"):
+        return f"""
+        **Classification Result:**
+        - **Label:** {result['label']}
+        - **Confidence:** {result['score']:.2%}
+        """
+    else:
+        return f"**Error:** {result.get('error', 'Unknown error')}"
+
+def api_classify(email_text: str) -> Dict[str, Any]:
+    """API endpoint function"""
+    return classify_single_email(email_text)
+
+def api_classify_batch(emails_json: str) -> str:
+    """Batch API endpoint function"""
+    try:
+        emails = json.loads(emails_json)
+        if not isinstance(emails, list):
+            return json.dumps({"error": "Input must be a JSON array of strings"})
+        
+        results = classify_batch_emails(emails)
+        return json.dumps(results, indent=2)
+    except json.JSONDecodeError:
+        return json.dumps({"error": "Invalid JSON format"})
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+# Load model on startup
+logger.info("Loading model...")
+model_loaded = load_model()
+
+if not model_loaded:
+    logger.warning("Model failed to load - using dummy responses")
+    def classify_single_email(email_text: str):
+        return {"label": "job", "score": 0.95, "success": True, "note": "Using dummy classifier"}
+
+# Create Gradio interface
+with gr.Blocks(title="Email Classifier", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 📧 Email Classification API")
+    gr.Markdown("Classify emails as job-related or other categories using a trained SetFit model.")
+    
+    with gr.Tab("Single Email Classification"):
+        with gr.Row():
+            with gr.Column():
+                email_input = gr.Textbox(
+                    label="Email Content",
+                    placeholder="Paste your email content here (subject + body)...",
+                    lines=8,
+                    max_lines=20
+                )
+                classify_btn = gr.Button("Classify Email", variant="primary")
+            
+            with gr.Column():
+                result_output = gr.Markdown(label="Classification Result")
+        
+        classify_btn.click(
+            fn=gradio_classify,
+            inputs=email_input,
+            outputs=result_output
+        )
+    
+    with gr.Tab("API Endpoints"):
+        gr.Markdown("""
+        ## API Usage
+        
+        ### Single Email Classification
+        **POST** `/api/classify`
+        ```json
+        {
+            "email_text": "Your email content here..."
+        }
+        ```
+        
+        ### Batch Email Classification  
+        **POST** `/api/classify-batch`
+        ```json
+        {
+            "emails": [
+                "Email 1 content...",
+                "Email 2 content...",
+                "Email 3 content..."
+            ]
+        }
+        ```
+        
+        ### Example Response
+        ```json
+        {
+            "label": "job",
+            "score": 0.9234,
+            "success": true
+        }
+        ```
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Test Single API")
+                api_input = gr.Textbox(label="Email Text", lines=4)
+                api_btn = gr.Button("Test API")
+                api_output = gr.JSON(label="API Response")
+                
+                api_btn.click(
+                    fn=api_classify,
+                    inputs=api_input,
+                    outputs=api_output
+                )
+            
+            with gr.Column():
+                gr.Markdown("### Test Batch API")
+                batch_input = gr.Textbox(
+                    label="JSON Array of Emails", 
+                    lines=6,
+                    placeholder='["Email 1 content", "Email 2 content"]'
+                )
+                batch_btn = gr.Button("Test Batch API")
+                batch_output = gr.Code(label="Batch API Response", language="json")
+                
+                batch_btn.click(
+                    fn=api_classify_batch,
+                    inputs=batch_input,
+                    outputs=batch_output
+                )
+    
+    with gr.Tab("Model Info"):
+        gr.Markdown(f"""
+        ### Model Information
+        - **Status:** {'✅ Loaded' if model_loaded else '❌ Failed to load'}
+        - **Model Type:** SetFit Email Classifier
+        - **Categories:** Job-related emails, Other emails
+        - **API Base URL:** `https://your-space-name.hf.space`
+        
+        ### Integration with Next.js
+        ```javascript
+        // Single email classification
+        const response = await fetch('https://your-space-name.hf.space/api/classify', {{
+            method: 'POST',
+            headers: {{ 'Content-Type': 'application/json' }},
+            body: JSON.stringify({{ email_text: emailContent }})
+        }});
+        const result = await response.json();
+        
+        // Batch classification
+        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 }})
+        }});
+        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
+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,
+        show_api=True,
+        share=False
+    )