Update main.py

main.py

@@ -1,10 +1,4 @@
-
 import os
-
-# Set safe Hugging Face cache directory before any transformers import
-os.environ["HF_HOME"] = "/tmp/hf-cache"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf-cache"
-
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
@@ -12,9 +6,15 @@ import torch
 
 app = FastAPI()
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")
-model = AutoModelForSequenceClassification.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")
+# Explicitly set cache directory in function (avoid HF_HOME confusion)
+tokenizer = AutoTokenizer.from_pretrained(
+    "cybersectony/phishing-email-detection-distilbert_v2.4.1",
+    cache_dir="/tmp"
+)
+model = AutoModelForSequenceClassification.from_pretrained(
+    "cybersectony/phishing-email-detection-distilbert_v2.4.1",
+    cache_dir="/tmp"
+)
 
 class EmailInput(BaseModel):
     text: str