Spaces:
Runtime error
Runtime error
File size: 2,467 Bytes
4df8249 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import torch
import transformers
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
def load_model(
base,
finetuned,
mode_cpu,
mode_mps,
mode_full_gpu,
mode_8bit,
mode_4bit,
# force_download_ckpt,
model_cls,
tokenizer_cls
):
if tokenizer_cls is None:
tokenizer_cls = AutoTokenizer
else:
tokenizer_cls = eval(tokenizer_cls)
if model_cls is None:
model_cls = AutoModelForCausalLM
else:
model_cls = eval(model_cls)
print(f"tokenizer_cls: {tokenizer_cls}")
print(f"model_cls: {model_cls}")
tokenizer = tokenizer_cls.from_pretrained(base)
tokenizer.padding_side = "left"
if mode_cpu:
print("cpu mode")
model = model_cls.from_pretrained(
base,
device_map={"": "cpu"},
use_safetensors=False
# low_cpu_mem_usage=True
)
if finetuned is not None and \
finetuned != "" and \
finetuned != "N/A":
model = PeftModel.from_pretrained(
model,
finetuned,
device_map={"": "cpu"}
# force_download=force_download_ckpt,
)
elif mode_mps:
print("mps mode")
model = model_cls.from_pretrained(
base,
device_map={"": "mps"},
torch_dtype=torch.float16,
use_safetensors=False
)
if finetuned is not None and \
finetuned != "" and \
finetuned != "N/A":
model = PeftModel.from_pretrained(
model,
finetuned,
torch_dtype=torch.float16,
device_map={"": "mps"}
# force_download=force_download_ckpt,
)
else:
print("gpu mode")
print(f"8bit = {mode_8bit}, 4bit = {mode_4bit}")
model = model_cls.from_pretrained(
base,
load_in_8bit=mode_8bit,
load_in_4bit=mode_4bit,
torch_dtype=torch.float16,
device_map="auto",
)
if finetuned is not None and \
finetuned != "" and \
finetuned != "N/A":
model = PeftModel.from_pretrained(
model,
finetuned,
# force_download=force_download_ckpt,
)
return model, tokenizer |