Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login
 import os
-from accelerate import init_empty_weights, infer_auto_device_map, dispatch_model
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
 # Retrieve the token from the environment variable
 hf_api_token = os.getenv("HF_API_TOKEN")
@@ -24,9 +24,13 @@ model_name = "meta-llama/Llama-2-7b-hf"
 with init_empty_weights():
     summarization_model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Infer device map and dispatch model
-device_map = infer_auto_device_map(summarization_model, max_memory={0: "14GiB", 1: "14GiB"})
-summarization_model = dispatch_model(summarization_model, device_map=device_map)
+# Load checkpoint and dispatch model
+summarization_model = load_checkpoint_and_dispatch(
+    summarization_model,
+    checkpoint=model_name,
+    device_map="auto",
+    dtype=torch.float16
+)
 
 summarization_tokenizer = AutoTokenizer.from_pretrained(model_name)