Spaces:
Runtime error
Runtime error
File size: 1,048 Bytes
88f55d9 |
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 |
import torch
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM
from optimum.bettertransformer import BetterTransformer
def load_model(
base,
finetuned,
mode_cpu,
mode_mps,
mode_full_gpu,
mode_8bit,
mode_4bit,
force_download_ckpt
):
tokenizer = AutoTokenizer.from_pretrained(
base, trust_remote_code=True)
tokenizer.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained(
base,
load_in_8bit=mode_8bit,
load_in_4bit=mode_4bit,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
)
if finetuned is not None and \
finetuned != "" and \
finetuned != "N/A":
model = PeftModel.from_pretrained(
model,
finetuned,
# force_download=force_download_ckpt,
trust_remote_code=True
)
model = model.merge_and_unload()
# model = BetterTransformer.transform(model)
model.to('cuda')
return model, tokenizer
|