Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import torch
 import logging
 import sys
+import os
 from accelerate import infer_auto_device_map, init_empty_weights
 
 # Configure logging
@@ -13,6 +14,12 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Get HuggingFace token from environment variable
+hf_token = os.environ.get('HUGGINGFACE_TOKEN')
+if not hf_token:
+    logger.error("HUGGINGFACE_TOKEN environment variable not set")
+    raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable")
+
 # Define the model name
 model_name = "meta-llama/Llama-2-7b-hf"
 
@@ -33,7 +40,7 @@ try:
     tokenizer = AutoTokenizer.from_pretrained(
         model_name,
         trust_remote_code=True,
-        token=True  # You'll need to set your HF token here
+        token=hf_token
     )
     tokenizer.pad_token = tokenizer.eos_token
     logger.info("Tokenizer loaded successfully")
@@ -44,7 +51,7 @@ try:
         model_name,
         torch_dtype=torch.float16 if device == "cuda" else torch.float32,
         trust_remote_code=True,
-        token=True,  # You'll need to set your HF token here
+        token=hf_token,
         device_map="auto",
         load_in_8bit=True
     )