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, token=huggingface_token)
     model = AutoModelForCausalLM.from_pretrained(
         model_id, 
         torch_dtype=dtype, 
         device_map="auto",
         quantization_config=quantization_config,
+        token=huggingface_token,
         low_cpu_mem_usage=True
     )
     return tokenizer, model
@@ -39,7 +45,6 @@ def moderate(user_input, assistant_response):
         )
     
     result = tokenizer.decode(output[0], skip_special_tokens=True)
-    
     result = result.split(assistant_response)[-1].strip()
     
     is_safe = "safe" in result.lower()