Spaces:

ohmreborn
/

llama-question-generation

Runtime error

ohmreborn commited on Jun 14, 2023

Commit

10711d5

•

1 Parent(s): 4c3d1f2

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -20,7 +20,8 @@ def main(
     model = LlamaForCausalLM.from_pretrained(
         base_model,
         load_in_8bit=load_8bit,
-        device_map={"": device},
         low_cpu_mem_usage=True
     )
     model = PeftModel.from_pretrained(
@@ -28,7 +29,6 @@ def main(
         lora_weights,
         device_map={"": device},
     )
-    model = model.merge_and_unload()
     model.config.pad_token_id = tokenizer.pad_token_id = 0

     model = LlamaForCausalLM.from_pretrained(
         base_model,
         load_in_8bit=load_8bit,
+        max_memory={"cpu":"15GiB"},
+        device_map="auto",
         low_cpu_mem_usage=True
     )
     model = PeftModel.from_pretrained(
         lora_weights,
         device_map={"": device},
     )
     model.config.pad_token_id = tokenizer.pad_token_id = 0