|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
from peft import PeftModel |
|
import torch |
|
|
|
import os |
|
import argparse |
|
|
|
def get_args(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--base_model_name_or_path", type=str, default="bigcode/large-model") |
|
parser.add_argument("--peft_model_path", type=str, default="/") |
|
parser.add_argument("--push_to_hub", action="store_true", default=True) |
|
|
|
return parser.parse_args() |
|
|
|
def main(): |
|
args = get_args() |
|
|
|
base_model = AutoModelForCausalLM.from_pretrained( |
|
args.base_model_name_or_path, |
|
return_dict=True, |
|
|
|
) |
|
|
|
|
|
model = PeftModel.from_pretrained(base_model, args.peft_model_path) |
|
model = model.merge_and_unload() |
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
tokenizer = AutoTokenizer.from_pretrained(args.base_model_name_or_path) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
model.save_pretrained(f"{args.peft_model_path}-merged") |
|
tokenizer.save_pretrained(f"{args.peft_model_path}-merged") |
|
print(f"Model saved to {args.peft_model_path}-merged") |
|
|
|
if __name__ == "__main__" : |
|
main() |