from transformers import AutoModelForCausalLM, BitsAndBytesConfig, HfArgumentParser, TrainingArguments
from peft import AutoPeftModelForCausalLM
from peft import PeftModel
from peft import LoraConfig, get_peft_model

#init
parser = HfArgumentParser(ScriptArguments)
script_args = parser.parse_args_into_dataclasses()[0]

# specific the model to load
# For GP-GPT small:
script_args.model_name = "meta-llama/Llama-2-7b"
script_args.peft_model_id = "./small/"

# For GP-GPT base:
script_args.model_name = "meta-llama/Meta-Llama-3.1-8B"
script_args.peft_model_id = "./base/"

# Cache model
model = AutoModelForCausalLM.from_pretrained(
        script_args.model_name,
        #quantization_config=quantization_config, # activate when using quantization setting
        device_map=device_map,
        torch_dtype=torch_dtype,
        use_auth_token=False,
    )

#load PEFT adapter
if script_args.peft_model_id is not None:
    peft_model_id = script_args.peft_model_id
    model = PeftModel.from_pretrained(model, peft_model_id)
    model = model.merge_and_unload()

IanL10/GP-GPT

Usage

IanL10/
GP-GPT