Update app.py

app.py

@@ -28,22 +28,14 @@ model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
 text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokenizer_gpt2)
 
 # Load the Llama-3 model and tokenizer once during startup
-tokenizer_llama = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B", token=hf_token)
-# Define the offload directory
-offload_dir = "./offload"  
-os.makedirs(offload_dir, exist_ok=True)  # Create the directory if it doesn't exist
-
-# Load the Llama model with disk offloading
-with init_empty_weights():
-    model_llama = AutoModelForCausalLM.from_pretrained(
-        "meta-llama/Meta-Llama-3-8B",
+tokenizer_llama = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B", token=hf_token)
+model_llama = AutoModelForCausalLM.from_pretrained(
+        "meta-llama/Meta-Llama-3.1-8B",
         torch_dtype='auto',
         device_map='auto',
         token=hf_token
     )
 
-# Offload the model to the specified directory
-disk_offload(model_llama, offload_dir)  # Pass the offload directory
 
 # Define your prompt template
 prompt_template = """\