Update app.py

app.py

@@ -7,6 +7,7 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, AutoMode
 from huggingface_hub import HfFolder
 from io import StringIO
 from tqdm import tqdm
+from accelerate import init_empty_weight
 
 # Access the Hugging Face API token from environment variables
 hf_token = os.getenv('HF_API_TOKEN')
@@ -27,12 +28,14 @@ text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokeniz
 
 # Load the Llama-3 model and tokenizer once during startup
 tokenizer_llama = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B", token=hf_token)
-model_llama = AutoModelForCausalLM.from_pretrained(
-    "meta-llama/Meta-Llama-3-8B",
-    torch_dtype='auto',
-    device_map='auto',
-    token=hf_token
-)
+
+with init_empty_weights():
+    model_llama = AutoModelForCausalLM.from_pretrained(
+        "meta-llama/Meta-Llama-3-8B",
+        torch_dtype='auto',
+        device_map='auto',
+        token=hf_token
+    )
 
 # Define your prompt template
 prompt_template = """\