mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-09-24 13:20:12 +08:00
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import os
|
|
# disable cuda
|
|
os.environ['CUDA_VISIBLE_DEVICES']="-1"
|
|
import torch
|
|
import argparse
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
model_options = dict(
|
|
device_map="auto",
|
|
# load_in_8bit=True, # not with cpu
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True
|
|
)
|
|
|
|
def main(model_repo, lora_repo = None, **download_options):
|
|
tokenizer = AutoTokenizer.from_pretrained(model_repo, **download_options)
|
|
model = AutoModelForCausalLM.from_pretrained(model_repo, **model_options, **download_options)
|
|
|
|
if lora_repo is not None:
|
|
# https://github.com/tloen/alpaca-lora/blob/main/generate.py#L40
|
|
from peft import PeftModel
|
|
model = PeftModel.from_pretrained(
|
|
model,
|
|
lora_repo,
|
|
torch_dtype=torch.float16,
|
|
device_map='auto',
|
|
**download_options
|
|
)
|
|
|
|
|
|
|
|
|
|
if __name__=="__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('model_repo', type=str)
|
|
parser.add_argument('-l', '--lora_repo', type=str, default=None, help='Name of the lora repo')
|
|
parser.add_argument('-f', '--force_download', type=str, default=None, help='Name of the lora repo')
|
|
parser.add_argument('-r', '--resume_download', type=str, default=None, help='Name of the lora repo')
|
|
args = parser.parse_args()
|
|
|
|
main(args.model_repo, args.lora_repo, force_download=args.force_download, resume_download=args.resume_download, low_cpu_mem_usage=True)
|