sandmanbuzz commited on
Commit
3a5a873
·
verified ·
1 Parent(s): 370c6dc

Handy Utility For Later

Browse files
Files changed (1) hide show
  1. loraize.py +53 -0
loraize.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # you have got to be shitting me
2
+ import huggingface_hub
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer
4
+ from peft import PeftModel
5
+ import torch
6
+
7
+ import os
8
+ import argparse
9
+
10
+ parser = argparse.ArgumentParser(
11
+ prog='loraize',
12
+ description='Apply one or more loras to a model, and then save it',
13
+ epilog='BOTTOM TEXT')
14
+ parser.add_argument(
15
+ 'model',
16
+ type=str,
17
+ help='path or HF name of a base model',
18
+ )
19
+ parser.add_argument(
20
+ 'lora',
21
+ type=str,
22
+ help='one or more LORAs to apply',
23
+ nargs='+')
24
+ parser.add_argument(
25
+ 'output_dir',
26
+ type=str,
27
+ help='output directory',
28
+ )
29
+ args = parser.parse_args()
30
+
31
+ print(f"Loading bassoon model:", args.model)
32
+ base_model = AutoModelForCausalLM.from_pretrained(
33
+ args.model,
34
+ return_dict=True,
35
+ torch_dtype=torch.bfloat16,
36
+ device_map="cpu",
37
+ )
38
+
39
+ for lora in args.lora:
40
+ print(f"Loading LORA: ",lora)
41
+ model = PeftModel.from_pretrained(
42
+ base_model,
43
+ lora,
44
+ device_map="cpu"
45
+ )
46
+ print(f"Good luck, bitches. Unloading.")
47
+ print("This gon' take a sec.")
48
+ model = model.merge_and_unload()
49
+ tokenizer = AutoTokenizer.from_pretrained(args.model)
50
+
51
+ model.save_pretrained(args.output_dir, safe_serialization=True, max_shard_size='10GB')
52
+ tokenizer.save_pretrained(args.output_dir)
53
+