fix offload_dir error
Browse files- utils/inference.py +8 -2
utils/inference.py
CHANGED
@@ -45,12 +45,18 @@ def load_tokenizer_and_model(base_model, adapter_model, load_8bit=True):
|
|
45 |
)
|
46 |
else:
|
47 |
model = LlamaForCausalLM.from_pretrained(
|
48 |
-
base_model,
|
|
|
|
|
|
|
|
|
49 |
)
|
50 |
if adapter_model is not None:
|
51 |
model = PeftModel.from_pretrained(
|
52 |
model,
|
53 |
-
adapter_model
|
|
|
|
|
54 |
)
|
55 |
|
56 |
model.eval()
|
|
|
45 |
)
|
46 |
else:
|
47 |
model = LlamaForCausalLM.from_pretrained(
|
48 |
+
base_model,
|
49 |
+
device_map={"": device},
|
50 |
+
low_cpu_mem_usage=True,
|
51 |
+
torch_dtype=torch.float16,
|
52 |
+
offload_folder="."
|
53 |
)
|
54 |
if adapter_model is not None:
|
55 |
model = PeftModel.from_pretrained(
|
56 |
model,
|
57 |
+
adapter_model,
|
58 |
+
torch_dtype=torch.float16,
|
59 |
+
offload_folder="."
|
60 |
)
|
61 |
|
62 |
model.eval()
|