Skip to content

Commit 7ba3481

Browse files
feat: allow specifying the chat template as base64 to avoid weird escaping and templating issues (#534)
test: create tests for chat_template_base64 field Signed-off-by: Harikrishnan Balagopal <harikrishmenon@gmail.com>
1 parent a84b716 commit 7ba3481

6 files changed

Lines changed: 171 additions & 0 deletions

File tree

tests/artifacts/predefined_data_configs/__init__.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,18 @@
4343
DATA_CONFIG_MULTITURN_CHAT_TOKENIZE_AND_MASKING_DATA_HANDLER = os.path.join(
4444
PREDEFINED_DATA_CONFIGS, "mt_data_granite_3_1B_tokenize_and_mask_handler.yaml"
4545
)
46+
DATA_CONFIG_VALID_BASE64_CHAT_TEMPLATE = os.path.join(
47+
PREDEFINED_DATA_CONFIGS,
48+
"granite_3_1b_valid_base64_data_handler.yaml",
49+
)
50+
DATA_CONFIG_INVALID_BASE64_CHAT_TEMPLATE = os.path.join(
51+
PREDEFINED_DATA_CONFIGS,
52+
"granite_3_1b_invalid_base64_data_handler.yaml",
53+
)
54+
GRANITE_3_1_B_CHAT_TEMPLATE = os.path.join(
55+
PREDEFINED_DATA_CONFIGS,
56+
"granite_3_1b_chat_template.txt",
57+
)
4658
DATA_CONFIG_YAML_STREAMING_INPUT_OUTPUT = os.path.join(
4759
PREDEFINED_DATA_CONFIGS, "tokenize_and_apply_input_masking_streaming.yaml"
4860
)
Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
{%- if messages[0]['role'] == 'system' %}
2+
{%- set system_message = messages[0]['content'] %}
3+
{%- set loop_messages = messages[1:] %}
4+
{%- else %}
5+
{%- set system_message = "Knowledge Cutoff Date: April 2024.\nToday's Date: " + strftime_now('%B %d, %Y') + ".\nYou are Granite, developed by IBM." %}
6+
{%- if tools and documents %}
7+
{%- set system_message = system_message + " You are a helpful AI assistant with access to the following tools. When a tool is required to answer the user's query, respond with <|tool_call|> followed by a JSON list of tools used. If a tool does not exist in the provided list of tools, notify the user that you do not have the ability to fulfill the request.\n\nWrite the response to the user's input by strictly aligning with the facts in the provided documents. If the information needed to answer the question is not available in the documents, inform the user that the question cannot be answered based on the available data." %}
8+
{%- elif tools %}
9+
{%- set system_message = system_message + " You are a helpful AI assistant with access to the following tools. When a tool is required to answer the user's query, respond with <|tool_call|> followed by a JSON list of tools used. If a tool does not exist in the provided list of tools, notify the user that you do not have the ability to fulfill the request." %}
10+
{%- elif documents %}
11+
{%- set system_message = system_message + " Write the response to the user's input by strictly aligning with the facts in the provided documents. If the information needed to answer the question is not available in the documents, inform the user that the question cannot be answered based on the available data." %}
12+
{%- else %}
13+
{%- set system_message = system_message + " You are a helpful AI assistant." %}
14+
{%- endif %}
15+
{%- if 'citations' in controls and documents %}
16+
{%- set system_message = system_message + '\n\nIn your response, use the symbols <co> and </co> to indicate when a fact comes from a document in the search result, e.g <co>0</co> for a fact from document 0. Afterwards, list all the citations with their corresponding documents in an ordered list.' %}
17+
{%- endif %}
18+
{%- if 'hallucinations' in controls and documents %}
19+
{%- set system_message = system_message + '\n\nFinally, after the response is written, include a numbered list of sentences from the response that are potentially hallucinated and not based in the documents.' %}
20+
{%- endif %}
21+
{%- set loop_messages = messages %}
22+
{%- endif %}
23+
{{- '<|start_of_role|>system<|end_of_role|>' + system_message + '<|end_of_text|>\n' }}
24+
{%- if tools %}
25+
{{- '<|start_of_role|>tools<|end_of_role|>' }}
26+
{{- tools | tojson(indent=4) }}
27+
{{- '<|end_of_text|>\n' }}
28+
{%- endif %}
29+
{%- if documents %}
30+
{{- '<|start_of_role|>documents<|end_of_role|>' }}
31+
{%- for document in documents %}
32+
{{- 'Document ' + loop.index0 | string + '\n' }}
33+
{{- document['text'] }}
34+
{%- if not loop.last %}
35+
{{- '\n\n'}}
36+
{%- endif%}
37+
{%- endfor %}
38+
{{- '<|end_of_text|>\n' }}
39+
{%- endif %}
40+
{%- for message in loop_messages %}
41+
{{- '<|start_of_role|>' + message['role'] + '<|end_of_role|>' + message['content'] + '<|end_of_text|>\n' }}
42+
{%- if loop.last and add_generation_prompt %}
43+
{{- '<|start_of_role|>assistant' }}
44+
{%- if controls %}
45+
{{- ' ' + controls | tojson()}}
46+
{%- endif %}
47+
{{- '<|end_of_role|>' }}
48+
{%- endif %}
49+
{%- endfor %}
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
dataprocessor:
2+
type: default
3+
chat_template_base64: 'dsaeyUtIGlmIG1lc3NhZ2VzWzBdWydyb2xlJ10gPT0gJ3N5c3RlbScgJX0KICAgIHslLSBzZXQgc3lzdGVtX21lc3NhZ2UgPSBtZXNzYWdlc1swXVsnY29udGVudCddICV9CiAgICB7JS0gc2V0IGxvb3BfbWVzc2FnZXMgPSBtZXNzYWdlc1sxOl0gJX0KeyUtIGVsc2UgJX0KICAgIHslLSBzZXQgc3lzdGVtX21lc3NhZ2UgPSAiS25vd2xlZGdlIEN1dG9mZiBEYXRlOiBBcHJpbCAyMDI0LlxuVG9kYXkncyBEYXRlOiAiICsgc3RyZnRpbWVfbm93KCclQiAlZCwgJVknKSArICIuXG5Zb3UgYXJlIEdyYW5pdGUsIGRldmVsb3BlZCBieSBJQk0uIiAlfQogICAgeyUtIGlmIHRvb2xzIGFuZCBkb2N1bWVudHMgJX0KICAgICAgICB7JS0gc2V0IHN5c3RlbV9tZXNzYWdlID0gc3lzdGVtX21lc3NhZ2UgKyAiIFlvdSBhcmUgYSBoZWxwZnVsIEFJIGFzc2lzdGFudCB3aXRoIGFjY2VzcyB0byB0aGUgZm9sbG93aW5nIHRvb2xzLiBXaGVuIGEgdG9vbCBpcyByZXF1aXJlZCB0byBhbnN3ZXIgdGhlIHVzZXIncyBxdWVyeSwgcmVzcG9uZCB3aXRoIDx8dG9vbF9jYWxsfD4gZm9sbG93ZWQgYnkgYSBKU09OIGxpc3Qgb2YgdG9vbHMgdXNlZC4gSWYgYSB0b29sIGRvZXMgbm90IGV4aXN0IGluIHRoZSBwcm92aWRlZCBsaXN0IG9mIHRvb2xzLCBub3RpZnkgdGhlIHVzZXIgdGhhdCB5b3UgZG8gbm90IGhhdmUgdGhlIGFiaWxpdHkgdG8gZnVsZmlsbCB0aGUgcmVxdWVzdC5cblxuV3JpdGUgdGhlIHJlc3BvbnNlIHRvIHRoZSB1c2VyJ3MgaW5wdXQgYnkgc3RyaWN0bHkgYWxpZ25pbmcgd2l0aCB0aGUgZmFjdHMgaW4gdGhlIHByb3ZpZGVkIGRvY3VtZW50cy4gSWYgdGhlIGluZm9ybWF0aW9uIG5lZWRlZCB0byBhbnN3ZXIgdGhlIHF1ZXN0aW9uIGlzIG5vdCBhdmFpbGFibGUgaW4gdGhlIGRvY3VtZW50cywgaW5mb3JtIHRoZSB1c2VyIHRoYXQgdGhlIHF1ZXN0aW9uIGNhbm5vdCBiZSBhbnN3ZXJlZCBiYXNlZCBvbiB0aGUgYXZhaWxhYmxlIGRhdGEuIiAlfQogICAgeyUtIGVsaWYgdG9vbHMgJX0KICAgICAgICB7JS0gc2V0IHN5c3RlbV9tZXNzYWdlID0gc3lzdGVtX21lc3NhZ2UgKyAiIFlvdSBhcmUgYSBoZWxwZnVsIEFJIGFzc2lzdGFudCB3aXRoIGFjY2VzcyB0byB0aGUgZm9sbG93aW5nIHRvb2xzLiBXaGVuIGEgdG9vbCBpcyByZXF1aXJlZCB0byBhbnN3ZXIgdGhlIHVzZXIncyBxdWVyeSwgcmVzcG9uZCB3aXRoIDx8dG9vbF9jYWxsfD4gZm9sbG93ZWQgYnkgYSBKU09OIGxpc3Qgb2YgdG9vbHMgdXNlZC4gSWYgYSB0b29sIGRvZXMgbm90IGV4aXN0IGluIHRoZSBwcm92aWRlZCBsaXN0IG9mIHRvb2xzLCBub3RpZnkgdGhlIHVzZXIgdGhhdCB5b3UgZG8gbm90IGhhdmUgdGhlIGFiaWxpdHkgdG8gZnVsZmlsbCB0aGUgcmVxdWVzdC4iICV9CiAgICB7JS0gZWxpZiBkb2N1bWVudHMgJX0KICAgICAgICB7JS0gc2V0IHN5c3RlbV9tZXNzYWdlID0gc3lzdGVtX21lc3NhZ2UgKyAiIFdyaXRlIHRoZSByZXNwb25zZSB0byB0aGUgdXNlcidzIGlucHV0IGJ5IHN0cmljdGx5IGFsaWduaW5nIHdpdGggdGhlIGZhY3RzIGluIHRoZSBwcm92aWRlZCBkb2N1bWVudHMuIElmIHRoZSBpbmZvcm1hdGlvbiBuZWVkZWQgdG8gYW5zd2VyIHRoZSBxdWVzdGlvbiBpcyBub3QgYXZhaWxhYmxlIGluIHRoZSBkb2N1bWVudHMsIGluZm9ybSB0aGUgdXNlciB0aGF0IHRoZSBxdWVzdGlvbiBjYW5ub3QgYmUgYW5zd2VyZWQgYmFzZWQgb24gdGhlIGF2YWlsYWJsZSBkYXRhLiIgJX0KICAgIHslLSBlbHNlICV9CiAgICAgICAgeyUtIHNldCBzeXN0ZW1fbWVzc2FnZSA9IHN5c3RlbV9tZXNzYWdlICsgIiBZb3UgYXJlIGEgaGVscGZ1bCBBSSBhc3Npc3RhbnQuIiAlfSAgICAKICAgIHslLSBlbmRpZiAlfQogICAgeyUtIGlmICdjaXRhdGlvbnMnIGluIGNvbnRyb2xzIGFuZCBkb2N1bWVudHMgJX0KICAgICAgICB7JS0gc2V0IHN5c3RlbV9tZXNzYWdlID0gc3lzdGVtX21lc3NhZ2UgKyAnXG5cbkluIHlvdXIgcmVzcG9uc2UsIHVzZSB0aGUgc3ltYm9scyA8Y28+IGFuZCA8L2NvPiB0byBpbmRpY2F0ZSB3aGVuIGEgZmFjdCBjb21lcyBmcm9tIGEgZG9jdW1lbnQgaW4gdGhlIHNlYXJjaCByZXN1bHQsIGUuZyA8Y28+MDwvY28+IGZvciBhIGZhY3QgZnJvbSBkb2N1bWVudCAwLiBBZnRlcndhcmRzLCBsaXN0IGFsbCB0aGUgY2l0YXRpb25zIHdpdGggdGhlaXIgY29ycmVzcG9uZGluZyBkb2N1bWVudHMgaW4gYW4gb3JkZXJlZCBsaXN0LicgJX0KICAgIHslLSBlbmRpZiAlfQogICAgeyUtIGlmICdoYWxsdWNpbmF0aW9ucycgaW4gY29udHJvbHMgYW5kIGRvY3VtZW50cyAlfQogICAgICAgIHslLSBzZXQgc3lzdGVtX21lc3NhZ2UgPSBzeXN0ZW1fbWVzc2FnZSArICdcblxuRmluYWxseSwgYWZ0ZXIgdGhlIHJlc3BvbnNlIGlzIHdyaXR0ZW4sIGluY2x1ZGUgYSBudW1iZXJlZCBsaXN0IG9mIHNlbnRlbmNlcyBmcm9tIHRoZSByZXNwb25zZSB0aGF0IGFyZSBwb3RlbnRpYWxseSBoYWxsdWNpbmF0ZWQgYW5kIG5vdCBiYXNlZCBpbiB0aGUgZG9jdW1lbnRzLicgJX0KICAgIHslLSBlbmRpZiAlfQogICAgeyUtIHNldCBsb29wX21lc3NhZ2VzID0gbWVzc2FnZXMgJX0KeyUtIGVuZGlmICV9Cnt7LSAnPHxzdGFydF9vZl9yb2xlfD5zeXN0ZW08fGVuZF9vZl9yb2xlfD4nICsgc3lzdGVtX21lc3NhZ2UgKyAnPHxlbmRfb2ZfdGV4dHw+XG4nIH19CnslLSBpZiB0b29scyAlfQogICAge3stICc8fHN0YXJ0X29mX3JvbGV8PnRvb2xzPHxlbmRfb2Zfcm9sZXw+JyB9fQogICAge3stIHRvb2xzIHwgdG9qc29uKGluZGVudD00KSB9fQogICAge3stICc8fGVuZF9vZl90ZXh0fD5cbicgfX0KeyUtIGVuZGlmICV9CnslLSBpZiBkb2N1bWVudHMgJX0KICAgIHt7LSAnPHxzdGFydF9vZl9yb2xlfD5kb2N1bWVudHM8fGVuZF9vZl9yb2xlfD4nIH19CiAgICB7JS0gZm9yIGRvY3VtZW50IGluIGRvY3VtZW50cyAlfQogICAgICAgIHt7LSAnRG9jdW1lbnQgJyArIGxvb3AuaW5kZXgwIHwgc3RyaW5nICsgJ1xuJyB9fQogICAgICAgIHt7LSBkb2N1bWVudFsndGV4dCddIH19CiAgICAgICAgeyUtIGlmIG5vdCBsb29wLmxhc3QgJX0KICAgICAgICAgICAge3stICdcblxuJ319CiAgICAgICAgeyUtIGVuZGlmJX0KICAgIHslLSBlbmRmb3IgJX0KICAgIHt7LSAnPHxlbmRfb2ZfdGV4dHw+XG4nIH19CnslLSBlbmRpZiAlfQp7JS0gZm9yIG1lc3NhZ2UgaW4gbG9vcF9tZXNzYWdlcyAlfQogICAge3stICc8fHN0YXJ0X29mX3JvbGV8PicgKyBtZXNzYWdlWydyb2xlJ10gKyAnPHxlbmRfb2Zfcm9sZXw+JyArIG1lc3NhZ2VbJ2NvbnRlbnQnXSArICc8fGVuZF9vZl90ZXh0fD5cbicgfX0KICAgIHslLSBpZiBsb29wLmxhc3QgYW5kIGFkZF9nZW5lcmF0aW9uX3Byb21wdCAlfQogICAgICAgIHt7LSAnPHxzdGFydF9vZl9yb2xlfD5hc3Npc3RhbnQnIH19CiAgICAgICAgICAgIHslLSBpZiBjb250cm9scyAlfQogICAgICAgICAgICAgICAge3stICcgJyArIGNvbnRyb2xzIHwgdG9qc29uKCl9fQogICAgICAgICAgICB7JS0gZW5kaWYgJX0KICAgICAgICB7ey0gJzx8ZW5kX29mX3JvbGV8PicgfX0KICAgIHslLSBlbmRpZiAlfQp7JS0gZW5kZm9yICV9Cg=='
4+
datasets:
5+
- name: dataset_1
6+
data_paths:
7+
- "FILE_PATH"
8+
data_handlers:
9+
- name: tokenize_and_apply_chat_template_with_masking
10+
arguments:
11+
remove_columns: all
12+
fn_kwargs:
13+
max_seq_length: 1024
14+
conversation_column: "messages"
15+
- name: dataset_2
16+
data_paths:
17+
- "FILE_PATH"
18+
data_handlers:
19+
- name: tokenize_and_apply_chat_template_with_masking
20+
arguments:
21+
remove_columns: all
22+
fn_kwargs:
23+
max_seq_length: 1024
24+
conversation_column: "messages"
25+
- name: dataset_3
26+
data_paths:
27+
- "FILE_PATH"
28+
data_handlers:
29+
- name: tokenize_and_apply_chat_template_with_masking
30+
arguments:
31+
remove_columns: all
32+
fn_kwargs:
33+
max_seq_length: 1024
34+
conversation_column: "messages"
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
dataprocessor:
2+
type: default
3+
chat_template_base64: 'eyUtIGlmIG1lc3NhZ2VzWzBdWydyb2xlJ10gPT0gJ3N5c3RlbScgJX0KICAgIHslLSBzZXQgc3lzdGVtX21lc3NhZ2UgPSBtZXNzYWdlc1swXVsnY29udGVudCddICV9CiAgICB7JS0gc2V0IGxvb3BfbWVzc2FnZXMgPSBtZXNzYWdlc1sxOl0gJX0KeyUtIGVsc2UgJX0KICAgIHslLSBzZXQgc3lzdGVtX21lc3NhZ2UgPSAiS25vd2xlZGdlIEN1dG9mZiBEYXRlOiBBcHJpbCAyMDI0LlxuVG9kYXkncyBEYXRlOiAiICsgc3RyZnRpbWVfbm93KCclQiAlZCwgJVknKSArICIuXG5Zb3UgYXJlIEdyYW5pdGUsIGRldmVsb3BlZCBieSBJQk0uIiAlfQogICAgeyUtIGlmIHRvb2xzIGFuZCBkb2N1bWVudHMgJX0KICAgICAgICB7JS0gc2V0IHN5c3RlbV9tZXNzYWdlID0gc3lzdGVtX21lc3NhZ2UgKyAiIFlvdSBhcmUgYSBoZWxwZnVsIEFJIGFzc2lzdGFudCB3aXRoIGFjY2VzcyB0byB0aGUgZm9sbG93aW5nIHRvb2xzLiBXaGVuIGEgdG9vbCBpcyByZXF1aXJlZCB0byBhbnN3ZXIgdGhlIHVzZXIncyBxdWVyeSwgcmVzcG9uZCB3aXRoIDx8dG9vbF9jYWxsfD4gZm9sbG93ZWQgYnkgYSBKU09OIGxpc3Qgb2YgdG9vbHMgdXNlZC4gSWYgYSB0b29sIGRvZXMgbm90IGV4aXN0IGluIHRoZSBwcm92aWRlZCBsaXN0IG9mIHRvb2xzLCBub3RpZnkgdGhlIHVzZXIgdGhhdCB5b3UgZG8gbm90IGhhdmUgdGhlIGFiaWxpdHkgdG8gZnVsZmlsbCB0aGUgcmVxdWVzdC5cblxuV3JpdGUgdGhlIHJlc3BvbnNlIHRvIHRoZSB1c2VyJ3MgaW5wdXQgYnkgc3RyaWN0bHkgYWxpZ25pbmcgd2l0aCB0aGUgZmFjdHMgaW4gdGhlIHByb3ZpZGVkIGRvY3VtZW50cy4gSWYgdGhlIGluZm9ybWF0aW9uIG5lZWRlZCB0byBhbnN3ZXIgdGhlIHF1ZXN0aW9uIGlzIG5vdCBhdmFpbGFibGUgaW4gdGhlIGRvY3VtZW50cywgaW5mb3JtIHRoZSB1c2VyIHRoYXQgdGhlIHF1ZXN0aW9uIGNhbm5vdCBiZSBhbnN3ZXJlZCBiYXNlZCBvbiB0aGUgYXZhaWxhYmxlIGRhdGEuIiAlfQogICAgeyUtIGVsaWYgdG9vbHMgJX0KICAgICAgICB7JS0gc2V0IHN5c3RlbV9tZXNzYWdlID0gc3lzdGVtX21lc3NhZ2UgKyAiIFlvdSBhcmUgYSBoZWxwZnVsIEFJIGFzc2lzdGFudCB3aXRoIGFjY2VzcyB0byB0aGUgZm9sbG93aW5nIHRvb2xzLiBXaGVuIGEgdG9vbCBpcyByZXF1aXJlZCB0byBhbnN3ZXIgdGhlIHVzZXIncyBxdWVyeSwgcmVzcG9uZCB3aXRoIDx8dG9vbF9jYWxsfD4gZm9sbG93ZWQgYnkgYSBKU09OIGxpc3Qgb2YgdG9vbHMgdXNlZC4gSWYgYSB0b29sIGRvZXMgbm90IGV4aXN0IGluIHRoZSBwcm92aWRlZCBsaXN0IG9mIHRvb2xzLCBub3RpZnkgdGhlIHVzZXIgdGhhdCB5b3UgZG8gbm90IGhhdmUgdGhlIGFiaWxpdHkgdG8gZnVsZmlsbCB0aGUgcmVxdWVzdC4iICV9CiAgICB7JS0gZWxpZiBkb2N1bWVudHMgJX0KICAgICAgICB7JS0gc2V0IHN5c3RlbV9tZXNzYWdlID0gc3lzdGVtX21lc3NhZ2UgKyAiIFdyaXRlIHRoZSByZXNwb25zZSB0byB0aGUgdXNlcidzIGlucHV0IGJ5IHN0cmljdGx5IGFsaWduaW5nIHdpdGggdGhlIGZhY3RzIGluIHRoZSBwcm92aWRlZCBkb2N1bWVudHMuIElmIHRoZSBpbmZvcm1hdGlvbiBuZWVkZWQgdG8gYW5zd2VyIHRoZSBxdWVzdGlvbiBpcyBub3QgYXZhaWxhYmxlIGluIHRoZSBkb2N1bWVudHMsIGluZm9ybSB0aGUgdXNlciB0aGF0IHRoZSBxdWVzdGlvbiBjYW5ub3QgYmUgYW5zd2VyZWQgYmFzZWQgb24gdGhlIGF2YWlsYWJsZSBkYXRhLiIgJX0KICAgIHslLSBlbHNlICV9CiAgICAgICAgeyUtIHNldCBzeXN0ZW1fbWVzc2FnZSA9IHN5c3RlbV9tZXNzYWdlICsgIiBZb3UgYXJlIGEgaGVscGZ1bCBBSSBhc3Npc3RhbnQuIiAlfSAgICAKICAgIHslLSBlbmRpZiAlfQogICAgeyUtIGlmICdjaXRhdGlvbnMnIGluIGNvbnRyb2xzIGFuZCBkb2N1bWVudHMgJX0KICAgICAgICB7JS0gc2V0IHN5c3RlbV9tZXNzYWdlID0gc3lzdGVtX21lc3NhZ2UgKyAnXG5cbkluIHlvdXIgcmVzcG9uc2UsIHVzZSB0aGUgc3ltYm9scyA8Y28+IGFuZCA8L2NvPiB0byBpbmRpY2F0ZSB3aGVuIGEgZmFjdCBjb21lcyBmcm9tIGEgZG9jdW1lbnQgaW4gdGhlIHNlYXJjaCByZXN1bHQsIGUuZyA8Y28+MDwvY28+IGZvciBhIGZhY3QgZnJvbSBkb2N1bWVudCAwLiBBZnRlcndhcmRzLCBsaXN0IGFsbCB0aGUgY2l0YXRpb25zIHdpdGggdGhlaXIgY29ycmVzcG9uZGluZyBkb2N1bWVudHMgaW4gYW4gb3JkZXJlZCBsaXN0LicgJX0KICAgIHslLSBlbmRpZiAlfQogICAgeyUtIGlmICdoYWxsdWNpbmF0aW9ucycgaW4gY29udHJvbHMgYW5kIGRvY3VtZW50cyAlfQogICAgICAgIHslLSBzZXQgc3lzdGVtX21lc3NhZ2UgPSBzeXN0ZW1fbWVzc2FnZSArICdcblxuRmluYWxseSwgYWZ0ZXIgdGhlIHJlc3BvbnNlIGlzIHdyaXR0ZW4sIGluY2x1ZGUgYSBudW1iZXJlZCBsaXN0IG9mIHNlbnRlbmNlcyBmcm9tIHRoZSByZXNwb25zZSB0aGF0IGFyZSBwb3RlbnRpYWxseSBoYWxsdWNpbmF0ZWQgYW5kIG5vdCBiYXNlZCBpbiB0aGUgZG9jdW1lbnRzLicgJX0KICAgIHslLSBlbmRpZiAlfQogICAgeyUtIHNldCBsb29wX21lc3NhZ2VzID0gbWVzc2FnZXMgJX0KeyUtIGVuZGlmICV9Cnt7LSAnPHxzdGFydF9vZl9yb2xlfD5zeXN0ZW08fGVuZF9vZl9yb2xlfD4nICsgc3lzdGVtX21lc3NhZ2UgKyAnPHxlbmRfb2ZfdGV4dHw+XG4nIH19CnslLSBpZiB0b29scyAlfQogICAge3stICc8fHN0YXJ0X29mX3JvbGV8PnRvb2xzPHxlbmRfb2Zfcm9sZXw+JyB9fQogICAge3stIHRvb2xzIHwgdG9qc29uKGluZGVudD00KSB9fQogICAge3stICc8fGVuZF9vZl90ZXh0fD5cbicgfX0KeyUtIGVuZGlmICV9CnslLSBpZiBkb2N1bWVudHMgJX0KICAgIHt7LSAnPHxzdGFydF9vZl9yb2xlfD5kb2N1bWVudHM8fGVuZF9vZl9yb2xlfD4nIH19CiAgICB7JS0gZm9yIGRvY3VtZW50IGluIGRvY3VtZW50cyAlfQogICAgICAgIHt7LSAnRG9jdW1lbnQgJyArIGxvb3AuaW5kZXgwIHwgc3RyaW5nICsgJ1xuJyB9fQogICAgICAgIHt7LSBkb2N1bWVudFsndGV4dCddIH19CiAgICAgICAgeyUtIGlmIG5vdCBsb29wLmxhc3QgJX0KICAgICAgICAgICAge3stICdcblxuJ319CiAgICAgICAgeyUtIGVuZGlmJX0KICAgIHslLSBlbmRmb3IgJX0KICAgIHt7LSAnPHxlbmRfb2ZfdGV4dHw+XG4nIH19CnslLSBlbmRpZiAlfQp7JS0gZm9yIG1lc3NhZ2UgaW4gbG9vcF9tZXNzYWdlcyAlfQogICAge3stICc8fHN0YXJ0X29mX3JvbGV8PicgKyBtZXNzYWdlWydyb2xlJ10gKyAnPHxlbmRfb2Zfcm9sZXw+JyArIG1lc3NhZ2VbJ2NvbnRlbnQnXSArICc8fGVuZF9vZl90ZXh0fD5cbicgfX0KICAgIHslLSBpZiBsb29wLmxhc3QgYW5kIGFkZF9nZW5lcmF0aW9uX3Byb21wdCAlfQogICAgICAgIHt7LSAnPHxzdGFydF9vZl9yb2xlfD5hc3Npc3RhbnQnIH19CiAgICAgICAgICAgIHslLSBpZiBjb250cm9scyAlfQogICAgICAgICAgICAgICAge3stICcgJyArIGNvbnRyb2xzIHwgdG9qc29uKCl9fQogICAgICAgICAgICB7JS0gZW5kaWYgJX0KICAgICAgICB7ey0gJzx8ZW5kX29mX3JvbGV8PicgfX0KICAgIHslLSBlbmRpZiAlfQp7JS0gZW5kZm9yICV9Cg=='
4+
datasets:
5+
- name: dataset_1
6+
data_paths:
7+
- "FILE_PATH"
8+
data_handlers:
9+
- name: tokenize_and_apply_chat_template_with_masking
10+
arguments:
11+
remove_columns: all
12+
fn_kwargs:
13+
max_seq_length: 1024
14+
conversation_column: "messages"
15+
- name: dataset_2
16+
data_paths:
17+
- "FILE_PATH"
18+
data_handlers:
19+
- name: tokenize_and_apply_chat_template_with_masking
20+
arguments:
21+
remove_columns: all
22+
fn_kwargs:
23+
max_seq_length: 1024
24+
conversation_column: "messages"
25+
- name: dataset_3
26+
data_paths:
27+
- "FILE_PATH"
28+
data_handlers:
29+
- name: tokenize_and_apply_chat_template_with_masking
30+
arguments:
31+
remove_columns: all
32+
fn_kwargs:
33+
max_seq_length: 1024
34+
conversation_column: "messages"

tests/test_sft_trainer.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
from scripts.run_inference import TunedCausalLM
3939
from tests.artifacts.predefined_data_configs import (
4040
DATA_CONFIG_DUPLICATE_COLUMNS,
41+
DATA_CONFIG_INVALID_BASE64_CHAT_TEMPLATE,
4142
DATA_CONFIG_MULTIPLE_DATASETS_SAMPLING_YAML,
4243
DATA_CONFIG_MULTITURN_CHAT_TOKENIZE_AND_MASKING_DATA_HANDLER,
4344
DATA_CONFIG_MULTITURN_DATA_YAML,
@@ -46,8 +47,10 @@
4647
DATA_CONFIG_SKIP_LARGE_TEXT_HANDLER,
4748
DATA_CONFIG_TOKENIZE_AND_APPLY_INPUT_MASKING_YAML,
4849
DATA_CONFIG_TOKENIZE_AND_TRAIN_WITH_HANDLER,
50+
DATA_CONFIG_VALID_BASE64_CHAT_TEMPLATE,
4951
DATA_CONFIG_YAML_STREAMING_INPUT_OUTPUT,
5052
DATA_CONFIG_YAML_STREAMING_PRETOKENIZED,
53+
GRANITE_3_1_B_CHAT_TEMPLATE,
5154
)
5255
from tests.artifacts.testdata import (
5356
CHAT_DATA_MULTI_TURN,
@@ -84,6 +87,7 @@
8487
DataHandlerConfig,
8588
DataPreProcessorConfig,
8689
DataSetConfig,
90+
load_and_validate_data_config,
8791
)
8892
from tuning.data.data_handlers import (
8993
DataHandler,
@@ -1391,6 +1395,25 @@ def test_run_chat_style_ft_using_dataconfig_for_chat_template(
13911395
assert 'Provide two rhyming words for the word "love"' in output_inference
13921396

13931397

1398+
def test_data_config_chat_template_as_base64():
1399+
"""Check that the chat_template specified as base64 is parsed correctly."""
1400+
expected_chat_template_path = GRANITE_3_1_B_CHAT_TEMPLATE
1401+
with open(expected_chat_template_path, "r", encoding="utf-8") as f:
1402+
expected_chat_template = f.read()
1403+
data_config_path = DATA_CONFIG_VALID_BASE64_CHAT_TEMPLATE
1404+
assert os.path.isfile(data_config_path)
1405+
data_config = load_and_validate_data_config(data_config_path)
1406+
parsed_chat_template = data_config.dataprocessor.chat_template
1407+
assert parsed_chat_template is not None, "the chat_template wasn't parsed correctly"
1408+
assert (
1409+
data_config.dataprocessor.chat_template == expected_chat_template
1410+
), "the chat_template wasn't parsed correctly"
1411+
# --------------------------------------------
1412+
with pytest.raises(ValueError):
1413+
data_config_path = DATA_CONFIG_INVALID_BASE64_CHAT_TEMPLATE
1414+
data_config = load_and_validate_data_config(data_config_path)
1415+
1416+
13941417
@pytest.mark.parametrize(
13951418
"data_args",
13961419
[

tuning/data/data_config.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
# Standard
16+
from base64 import b64decode
1617
from dataclasses import dataclass
1718
from typing import Dict, List, Optional
1819
import logging
@@ -153,6 +154,24 @@ def _validate_dataprocessor_config(dataprocessor_config) -> DataPreProcessorConf
153154
chat_template = kwargs["chat_template"]
154155
assert isinstance(chat_template, str), "chat_template should be a string"
155156
c.chat_template = chat_template
157+
elif "chat_template_base64" in kwargs:
158+
chat_template_base64 = kwargs["chat_template_base64"]
159+
assert isinstance(
160+
chat_template_base64, str
161+
), "chat_template_base64 should be a string"
162+
logger.warning(
163+
"You are using the 'chat_template_base64' field. "
164+
+ "Please use the 'chat_template' field instead for better readability."
165+
)
166+
try:
167+
chat_template_bytes = b64decode(chat_template_base64)
168+
chat_template = chat_template_bytes.decode("utf-8")
169+
c.chat_template = chat_template
170+
except Exception as e:
171+
raise ValueError(
172+
"You passed the 'chat_template_base64' field which failed during decoding."
173+
+ "Please check it or use a decoded chat template with the 'chat_template' field."
174+
) from e
156175
return c
157176

158177

0 commit comments

Comments
 (0)