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, load_checkpoint_and_dispatch
+from accelerate import init_empty_weights, infer_auto_device_map, dispatch_model
 
 # Retrieve the token from the environment variable
 hf_api_token = os.getenv("HF_API_TOKEN")
@@ -24,12 +24,9 @@ model_name = "meta-llama/Llama-2-7b-hf"
 with init_empty_weights():
     summarization_model = AutoModelForCausalLM.from_pretrained(model_name)
 
-summarization_model = load_checkpoint_and_dispatch(
-    summarization_model,
-    checkpoint_path=model_name,
-    device_map="auto",
-    dtype=torch.float16
-)
+# 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)
 
 summarization_tokenizer = AutoTokenizer.from_pretrained(model_name)