tianyang commited on
Commit
3a5a33e
·
1 Parent(s): 1bcf3fd

fix offload_dir error

Browse files
Files changed (1) hide show
  1. 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, device_map={"": device}, low_cpu_mem_usage=True, torch_dtype=torch.float16
 
 
 
 
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()