Merge branch 'main' of hf.co:spaces/iatechaways/mail_bot

model_loader.py

@@ -3,6 +3,7 @@ import os
 import torch
 import gc
 import json
+from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from peft import PeftModel
 
@@ -10,6 +11,9 @@ from peft import PeftModel
 BASE_MODEL = "meta-llama/Meta-Llama-3.2-3B"   # or local path if you prefer
 ADAPTER = "iatechaways/crm_model"            # your LoRA adapter on HF
 
+# Get HF token from environment
+HF_TOKEN = os.getenv("HF_TOKEN")
+
 # NOTE: For ZeroGPU we LOAD on CPU and move to GPU only during inference.
 def load_model_cpu():
     """
@@ -17,7 +21,7 @@ def load_model_cpu():
     Return a dict that holds references to model and tokenizer.
     """
     # load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True, token=HF_TOKEN)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "right"
@@ -27,7 +31,8 @@ def load_model_cpu():
         BASE_MODEL,
         device_map={"": "cpu"},
         torch_dtype=torch.float32,
-        low_cpu_mem_usage=True
+        low_cpu_mem_usage=True,
+                token=HF_TOKEN
     )
 
     model = PeftModel.from_pretrained(
@@ -36,7 +41,8 @@ def load_model_cpu():
         device_map={"": "cpu"},
         torch_dtype=torch.float32,
         is_trainable=False,
-        low_cpu_mem_usage=True
+        low_cpu_mem_usage=True,
+                token=HF_TOKEN
     )
 
     # don't create pipeline here; we'll call generate with model directly