Update app.py

app.py

@@ -4,51 +4,56 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel, PeftConfig
 from fastapi.middleware.cors import CORSMiddleware
 import torch
-from dotenv import load_dotenv
-import os
-
-
-load_dotenv()
-hf_token = os.getenv("HF_TOKEN")
-
 
 app = FastAPI()
 
+# Allow CORS for all origins (adjust this in production)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], 
+    allow_origins=["*"],  
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-adapter_path = "./checkpoint-711"  
+# Path to your HF Hub repo with full model + adapter
+adapter_path = "imnim/multi-label-email-classifier"
 
 try:
-    peft_config = PeftConfig.from_pretrained(adapter_path)
-
-  
+    # Load PEFT config to get base model path
+    peft_config = PeftConfig.from_pretrained(adapter_path, use_auth_token=True)
+    
+    # Load base model and tokenizer with HF auth token
     base_model = AutoModelForCausalLM.from_pretrained(
         peft_config.base_model_name_or_path,
         torch_dtype=torch.float32,
-        device_map={"": "cpu"}
+        device_map={"": "cpu"},
+        use_auth_token=True
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        peft_config.base_model_name_or_path,
+        use_auth_token=True
     )
-    tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
-
 
-    model = PeftModel.from_pretrained(base_model, adapter_path, device_map={"": "cpu"})
+    # Load adapter with HF auth token
+    model = PeftModel.from_pretrained(
+        base_model, adapter_path,
+        device_map={"": "cpu"},
+        use_auth_token=True
+    )
 
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    # Setup text-generation pipeline
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
 
 except Exception as e:
     raise RuntimeError(f"❌ Failed to load model + adapter: {str(e)}")
 
-# === Request Schema ===
+# Request schema
 class EmailInput(BaseModel):
     subject: str
     body: str
 
-# === Endpoint ===
+# POST /classify endpoint
 @app.post("/classify")
 async def classify_email(data: EmailInput):
     prompt = f"""### Subject:\n{data.subject}\n\n### Body:\n{data.body}\n\n### Labels:"""