mirror of
https://github.com/zenlm/engine.git
synced 2026-10-30 15:00:05 +00:00
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
from peft.tuners import lora
|
|
from transformers import AutoModelForCausalLM # type: ignore
|
|
import json
|
|
from peft.tuners.lora.config import LoraConfig
|
|
|
|
model_id = input("Enter the base model ID: ")
|
|
target_modules_in = input("Enter the target modules as a comma delimited list: ")
|
|
target_modules = target_modules_in.split(",")
|
|
target_modules = [x for x in target_modules if len(x) > 0]
|
|
target_modules = [x.strip() for x in target_modules]
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(model_id)
|
|
lora_config = LoraConfig(target_modules=target_modules, init_lora_weights=False)
|
|
|
|
model.add_adapter(lora_config, "default")
|
|
|
|
total_swapped = 0
|
|
loras = {}
|
|
for n, module in model.named_modules():
|
|
if isinstance(module, lora.Linear):
|
|
loras[n.split("lora_A.")[0]] = total_swapped
|
|
total_swapped += 1
|
|
elif isinstance(module, lora.Embedding):
|
|
loras[n.split("lora_embedding_A.")[0]] = total_swapped
|
|
total_swapped += 1
|
|
elif isinstance(module, lora.Conv2d):
|
|
loras[n.split("lora_A.")[0]] = total_swapped
|
|
total_swapped += 1
|
|
|
|
adapters_in = input(
|
|
"Enter a comma delimited list of adapter names as they were specified when training: "
|
|
)
|
|
adapters = adapters_in.split(",")
|
|
adapters = [x for x in adapters if len(x) > 0]
|
|
adapters = [x.strip() for x in adapters]
|
|
|
|
out = {"order": adapters, "layers": loras, "base_model_id": model_id}
|
|
|
|
outfile = input("Enter output file: ")
|
|
with open(outfile, "w") as f:
|
|
f.write(json.dumps(out))
|