Update app.py

app.py

@@ -4,6 +4,12 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoModelForCausalLM
 from io import StringIO
 import os
 
+# Access the Hugging Face API token from environment variables
+hf_token = os.getenv('HF_API_TOKEN')
+
+if not hf_token:
+    raise ValueError("Hugging Face API token is not set. Please set the HF_API_TOKEN environment variable.")
+HfFolder.save_token(hf_token)
 # Load the GPT-2 tokenizer and model
 tokenizer_gpt2 = GPT2Tokenizer.from_pretrained('gpt2')
 model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
@@ -11,7 +17,7 @@ model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
 # Load the Llama3 model in sharded mode
 model_name = "meta-llama/Meta-Llama-3.1-8B"
 try:
-    model_llama = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")  # use device_map for automatic sharding
+    model_llama = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", token = hf_token)  # use device_map for automatic sharding
 except OSError as e:
     print(f"Error loading model: {e}")