Skip to content

Commit 662a84d

Browse files
wangpengfei1013zhanghao.smooth
authored andcommitted
add support for iei yuan2.0 (https://huggingface.co/IEITYuan) (lm-sys#2919)
1 parent 87ed32d commit 662a84d

File tree

4 files changed

+41
-0
lines changed

4 files changed

+41
-0
lines changed

docs/model_support.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
- [WizardLM/WizardLM-13B-V1.0](https://huggingface.co/WizardLM/WizardLM-13B-V1.0)
6262
- [WizardLM/WizardCoder-15B-V1.0](https://huggingface.co/WizardLM/WizardCoder-15B-V1.0)
6363
- [Xwin-LM/Xwin-LM-7B-V0.1](https://huggingface.co/Xwin-LM/Xwin-LM-70B-V0.1)
64+
- [IEITYuan/Yuan2-2B/51B/102B-hf](https://huggingface.co/IEITYuan)
6465
- Any [EleutherAI](https://huggingface.co/EleutherAI) pythia model such as [pythia-6.9b](https://huggingface.co/EleutherAI/pythia-6.9b)
6566
- Any [Peft](https://github.com/huggingface/peft) adapter trained on top of a
6667
model above. To activate, must have `peft` in the model path. Note: If

fastchat/conversation.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1440,6 +1440,22 @@ def get_conv_template(name: str) -> Conversation:
14401440
)
14411441

14421442

1443+
#yuan 2.0 template
1444+
#reference:https://github.com/IEIT-Yuan/Yuan-2.0
1445+
#reference:https://huggingface.co/IEITYuan
1446+
register_conv_template(
1447+
Conversation(
1448+
name="yuan",
1449+
system_template="",
1450+
roles=("", ""),
1451+
sep_style=SeparatorStyle.NO_COLON_SINGLE,
1452+
sep="<sep>",
1453+
stop_str="<eod>",
1454+
1455+
)
1456+
)
1457+
1458+
14431459
if __name__ == "__main__":
14441460
from fastchat.conversation import get_conv_template
14451461

fastchat/model/model_adapter.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2159,6 +2159,22 @@ def match(self, model_path: str):
21592159
def get_default_conv_template(self, model_path: str) -> Conversation:
21602160
return get_conv_template("solar")
21612161

2162+
class Yuan2Adapter(BaseModelAdapter):
2163+
"""The model adapter for Yuan """
2164+
2165+
def match(self, model_path: str):
2166+
return "yuan" in model_path.lower()
2167+
2168+
def load_model(self, model_path: str, from_pretrained_kwargs: dict):
2169+
model, tokenizer = super().load_model(model_path, from_pretrained_kwargs)
2170+
tokenizer.add_tokens(
2171+
['<sep>', '<pad>', '<mask>', '<predict>', '<FIM_SUFFIX>', '<FIM_PREFIX>', '<FIM_MIDDLE>', '<commit_before>',
2172+
'<commit_msg>', '<commit_after>', '<jupyter_start>', '<jupyter_text>', '<jupyter_code>',
2173+
'<jupyter_output>', '<empty_output>'], special_tokens=True)
2174+
return model, tokenizer
2175+
2176+
def get_default_conv_template(self, model_path: str) -> Conversation:
2177+
return get_conv_template("yuan")
21622178

21632179
# Note: the registration order matters.
21642180
# The one registered earlier has a higher matching priority.
@@ -2246,6 +2262,7 @@ def get_default_conv_template(self, model_path: str) -> Conversation:
22462262
register_model_adapter(MetaMathAdapter)
22472263
register_model_adapter(BagelAdapter)
22482264
register_model_adapter(SolarAdapter)
2265+
register_model_adapter(Yuan2Adapter)
22492266

22502267
# After all adapters, try the default base adapter.
22512268
register_model_adapter(BaseModelAdapter)

fastchat/model/model_registry.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -543,3 +543,10 @@ def get_model_info(name: str) -> ModelInfo:
543543
"https://huggingface.co/meta-math",
544544
"MetaMath is a finetune of Llama2 on [MetaMathQA](https://huggingface.co/datasets/meta-math/MetaMathQA) that specializes in mathematical reasoning.",
545545
)
546+
547+
register_model_info(
548+
["Yuan2-2B-hf","Yuan2-51B-hf","Yuan2-102B-hf"],
549+
"IEIYuan",
550+
"https://huggingface.co/IEITYuan",
551+
"Yuan2 is a Basemodel developed by IEI.",
552+
)

0 commit comments

Comments
 (0)