Update app.py

app.py

@@ -25,7 +25,7 @@ from IPython.display import Markdown, display
 peft_model_id = "./"
 config = PeftConfig.from_pretrained(peft_model_id)
 
-bnb_config = BitsAndBytesConfig(
+quantization_config = BitsAndBytesConfig(
     load_in_8bit_fp32_cpu_offload=True,
     bnb_8bit_use_fp16=False,
     load_in_4bit=True,
@@ -34,11 +34,19 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
+device_map = {
+    "transformer.word_embeddings": "cpu",
+    "transformer.word_embeddings_layernorm": "cpu",
+    "lm_head": "cpu",
+    "transformer.h": "cpu",
+    "transformer.ln_f": "cpu",
+}
+
 model = AutoModelForCausalLM.from_pretrained(
     config.base_model_name_or_path,
     return_dict=True,
     quantization_config=bnb_config,
-    device_map=None,
+    device_map=device_map,
     trust_remote_code=True,
 )
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)