Update app.py

app.py

@@ -1,8 +1,13 @@
+import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import gradio as gr
 import spaces
 
+huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
+if not huggingface_token:
+    raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
+
 model_id = "meta-llama/Llama-Guard-3-8B-INT8"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16
@@ -10,12 +15,13 @@ dtype = torch.bfloat16
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=huggingface_token)
     model = AutoModelForCausalLM.from_pretrained(
         model_id, 
         torch_dtype=dtype, 
         device_map=device, 
-        quantization_config=quantization_config
+        quantization_config=quantization_config,
+        use_auth_token=huggingface_token
     )
     return tokenizer, model