@@ -1078,11 +1078,48 @@ def _set_vocab_rwkv_world(self):
10781078 tokens .append (f"[PAD{ i } ]" .encode ("utf-8" ))
10791079 toktypes .append (gguf .TokenType .UNUSED )
10801080
1081+ # hack: rwkv doesn't have a consistent hf implementation,
1082+ # so we still hardcode the template here for now
1083+ template = """{%- if not add_generation_prompt is defined -%}
1084+ {%- set add_generation_prompt = true -%}
1085+ {%- endif -%}
1086+ {%- set ns = namespace(system_prompt='') -%}
1087+ {%- for message in messages -%}
1088+ {%- if message['role'] == 'system' -%}
1089+ {%- set ns.system_prompt = message['content'] -%}
1090+ {%- endif -%}
1091+ {%- endfor -%}
1092+ {{bos_token}}
1093+ {%- if ns.system_prompt != '' -%}
1094+ {{- 'System: ' + ns.system_prompt + '\\ n\\ n' -}}
1095+ {%- endif -%}
1096+ {%- for message in messages -%}
1097+ {%- if message['role'] == 'user' -%}
1098+ {{- 'User: ' + message['content']|trim + '\\ n\\ n' -}}
1099+ {%- endif -%}
1100+ {%- if message['role'] == 'assistant' and message['content'] is not none -%}
1101+ {%- set content = message['content'] -%}
1102+ {%- if '</think>' in content -%}
1103+ {%- set content = content.split('</think>')[-1] -%}
1104+ {%- endif -%}
1105+ {{- 'Assistant: ' + content|trim + '\\ n\\ n' -}}
1106+ {%- endif -%}
1107+ {%- endfor -%}
1108+ {%- if add_generation_prompt -%}
1109+ {{- 'Assistant:' -}}
1110+ {%- if enable_thinking is defined and enable_thinking is false %}
1111+ {{- ' <think>\n </think>' }}
1112+ {%- endif %}
1113+ {%- if enable_thinking is defined and enable_thinking is true %}
1114+ {{- ' <think>' }}
1115+ {%- endif %}
1116+ {%- endif -%}"""
1117+
10811118 self .gguf_writer .add_tokenizer_model ("rwkv" )
10821119 self .gguf_writer .add_token_list (tokens )
10831120 self .gguf_writer .add_token_types (toktypes )
10841121 special_vocab = gguf .SpecialVocab (self .dir_model , load_merges = False )
1085- special_vocab .chat_template = "rwkv-world"
1122+ special_vocab .chat_template = template
10861123 # hack: Add '\n\n' as the EOT token to make it chat normally
10871124 special_vocab ._set_special_token ("eot" , 261 )
10881125 # hack: Override these as they have already been set (incorrectly)
0 commit comments