Update main.py

main.py

@@ -2,6 +2,10 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+import os
+
+# Set cache directory to avoid permission errors
+os.environ["HF_HOME"] = "/tmp/hf-cache"
 
 app = FastAPI()
 
@@ -9,7 +13,6 @@ app = FastAPI()
 tokenizer = AutoTokenizer.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")
 model = AutoModelForSequenceClassification.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")
 
-# Define input schema
 class EmailInput(BaseModel):
     text: str
 
@@ -19,7 +22,7 @@ def predict(input: EmailInput):
     with torch.no_grad():
         outputs = model(**inputs)
         predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    
+
     probs = predictions[0].tolist()
     labels = {
         "legitimate_email": probs[0],
@@ -28,7 +31,7 @@ def predict(input: EmailInput):
         "phishing_url": probs[3]
     }
     max_label = max(labels.items(), key=lambda x: x[1])
-    
+
     return {
         "prediction": max_label[0],
         "confidence": round(max_label[1], 4),