Views
No views yet
1# Import Library
2from transformers import AutoTokenizer, AutoModelForCausalLM
3import torch
4
5# Load the target model to be applied skill
6tokenizer = AutoTokenizer.from_pretrained(
7 "tokyotech-llm/Swallow-7b-hf"
8)
9model = AutoModelForCausalLM.from_pretrained(
10 "tokyotech-llm/Swallow-7b-hf",
11 torch_dtype=torch.bfloat16,
12 device_map="auto",
13)
14
15# Load SkillTree
16skill_tree = AutoModelForCausalLM.from_pretrained(
17 "HachiML/SkillTree-llama2-7b-hf-Code",
18 torch_dtype=torch.bfloat16,
19 device_map="auto",
20)
21
22# Apply the skill to the target model
23def apply_skill(model, skill_tree):
24 # excluded object
25 skip_layers = ["model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"]
26 # apply skill
27 for k, v in model.state_dict().items():
28 # layernorm is also excluded
29 if (k in skip_layers) or ("layernorm" in k):
30 continue
31 vector = skill_tree.state_dict()[k]
32 new_v = v + vector.to(v.device)
33 v.copy_(new_v)
34 return model
35
36model = apply_skill(model, skill_tree)
37
38# Push to hub
39model_name = "HachiML/Swallow-7b-hf-CodeSkill"
40tokenizer.save_pretrained(f"./models/{model_name}", repo_id=model_name, push_to_hub=True)
41model.save_pretrained(f"./models/{model_name}", repo_id=model_name, push_to_hub=True)