Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import spaces
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
@@ -5,6 +6,11 @@ import torch
 import logging
 import sys
 from accelerate import infer_auto_device_map, init_empty_weights
+from huggingface_hub import login
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
 
 # Configure logging
 logging.basicConfig(
@@ -13,6 +19,20 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Get HuggingFace token from environment variable
+hf_token = os.getenv('HUGGINGFACE_TOKEN')
+if not hf_token:
+    logger.error("HUGGINGFACE_TOKEN environment variable not found")
+    raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable")
+
+# Login to Hugging Face
+try:
+    login(token=hf_token)
+    logger.info("Successfully logged in to Hugging Face")
+except Exception as e:
+    logger.error(f"Failed to login to Hugging Face: {str(e)}")
+    raise
+
 # Define the model name
 model_name = "meta-llama/Llama-2-7b-hf"
 
@@ -32,8 +52,7 @@ try:
     logger.info("Loading tokenizer...")
     tokenizer = AutoTokenizer.from_pretrained(
         model_name,
-        trust_remote_code=True,
-        use_auth_token=True
+        trust_remote_code=True
     )
     logger.info("Tokenizer loaded successfully")
 
@@ -43,7 +62,6 @@ try:
         model_name,
         torch_dtype=torch.float16 if device == "cuda" else torch.float32,
         trust_remote_code=True,
-        use_auth_token=True,
         load_in_8bit=True,
         device_map="auto"
     )