File size: 536 Bytes
3ac5300 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from peft import PeftModel
from transformers import LlamaTokenizer, LlamaForCausalLM
def load_model(
base="decapoda-research/llama-7b-hf",
finetuned="tloen/alpaca-lora-7b",
):
tokenizer = LlamaTokenizer.from_pretrained(base)
tokenizer.pad_token_id = 0
tokenizer.padding_side = "left"
model = LlamaForCausalLM.from_pretrained(
base,
load_in_8bit=True,
device_map="auto",
)
model = PeftModel.from_pretrained(model, finetuned, device_map={'': 0})
return model, tokenizer
|