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