tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

imrpove llama check and fix safetensors file check

Browse Source
This commit is contained in:
Wing Lian
2023-04-17 23:49:21 -04:00
parent e1076430ff
commit 69164da079
1 changed files with 2 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
scripts/finetune.py
View File

@@ -85,14 +85,12 @@ def load_model(base_model, base_model_config, model_type, tokenizer_type, cfg, a
raise e raise e
try: try:
if cfg.load_4bit and "llama" in base_model: if cfg.load_4bit and "llama" in base_model or "llama" in cfg.model_type.lower():
from alpaca_lora_4bit.autograd_4bit import load_llama_model_4bit_low_ram from alpaca_lora_4bit.autograd_4bit import load_llama_model_4bit_low_ram
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
cache_model_path = Path(snapshot_download(base_model)) cache_model_path = Path(snapshot_download(base_model))
# TODO search .glob for a .pt, .safetensor, or .bin files = list(cache_model_path.glob('*.pt')) + list(cache_model_path.glob('*.safetensors')) + list(cache_model_path.glob('*.bin'))
cache_model_path.glob("*.pt")
files = list(cache_model_path.glob('*.pt')) + list(cache_model_path.glob('*.safetensor')) + list(cache_model_path.glob('*.bin'))
if len(files) > 0: if len(files) > 0:
model_path = str(files[0]) model_path = str(files[0])
else: else: