Update app.py

app.py

@@ -1,11 +1,15 @@
+import os
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the tokenizer and model
+# Retrieve the Hugging Face token from environment variables
+hf_token = os.getenv("HUGGINGFACE_TOKEN")
+
+# Load the tokenizer and model with the token
 model_id = "meta-llama/Llama-3.2-3B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto", use_auth_token=hf_token)
 
 # Define the prediction function
 def generate_text(prompt):