diff --git a/tests/python/test_to_sharegpt_optional_none.py b/tests/python/test_to_sharegpt_optional_none.py new file mode 100644 index 0000000000..05fcb8a6a8 --- /dev/null +++ b/tests/python/test_to_sharegpt_optional_none.py @@ -0,0 +1,97 @@ +import ast +import re +from pathlib import Path + + +def _load_formatter_builders(): + # Extract _parse_combined_prompt and _create_formatter without importing + # unsloth (importing unsloth needs unsloth_zoo / a GPU). Both are pure + # Python and only use the `re` module. + source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + wanted = {"_parse_combined_prompt", "_create_formatter"} + funcs = [ + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted + ] + namespace = {"re": re} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace["_parse_combined_prompt"], namespace["_create_formatter"] + + +class _StubDataset: + def __init__(self, column_names): + self.column_names = column_names + + +def _render(merged_prompt, columns, batch): + parse, create = _load_formatter_builders() + possible_columns, final_optional_prompts = parse(merged_prompt, _StubDataset(columns)) + processor = create(possible_columns, final_optional_prompts, "text") + return processor(batch)["text"] + + +def test_optional_block_missing_second_column_does_not_render_none(): + # A [[...]] block may reference several columns; only the first gates the + # block. A later column that is None must not render as the literal "None". + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": [None]}, + ) + assert out[0] == "Location: Paris, end" + assert "None" not in out[0] + + +def test_optional_block_all_columns_present_unchanged(): + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": ["France"]}, + ) + assert out[0] == "Location: Paris, France end" + + +def test_optional_block_gating_column_empty_is_dropped(): + # When the gating (first) column is empty the whole block is omitted; this + # behaviour is unchanged by the None coercion. + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": [""], "country": ["France"]}, + ) + assert out[0] == "Location: end" + + +def test_single_column_optional_block_gated_out_on_none(): + # Single-column blocks were already gated correctly (the sole column is the + # gate); confirm they stay unaffected. + merged_prompt = "Name: [[{name}]]!" + out = _render(merged_prompt, ["name"], {"name": [None, "Bob"]}) + assert out == ["Name: !", "Name: Bob!"] + + +def test_required_column_none_does_not_render_none(): + # A required (non-[[...]]) column that is None must not render as the + # literal "None" either; coercion happens at the row source, so both the + # required and optional branches are covered. + merged_prompt = "Location: {city}, {country} end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": [None]}, + ) + assert out[0] == "Location: Paris, end" + assert "None" not in out[0] + + +def test_optional_block_falsy_but_present_gating_value_still_renders(): + # The gate keeps a block whenever the first column is not "". A falsy but + # real value (0) must not be treated as absent, so the block still renders. + merged_prompt = "Count: [[{n}]]!" + out = _render(merged_prompt, ["n"], {"n": [0]}) + assert out[0] == "Count: 0!" diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 169b2dbd0e..dd1e433471 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2185,7 +2185,16 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name texts = [] for row_idx in range(n_rows): - row_values = {column: examples[column][row_idx] for column in columns} + # Coerce missing (None) columns to "" so they do not render as the + # literal string "None" in the emitted text. In a [[...]] block only + # the first column gates the block, so a later column can still be + # None here; required columns can be None too. Coercing at the source + # covers both; since None is now "", the gate below only needs to + # test for "" (an empty first column still drops the block). + row_values = { + column: ("" if (value := examples[column][row_idx]) is None else value) + for column in columns + } formatter_values = {} for formatter_template in formatter_templates: @@ -2196,7 +2205,7 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name continue _, optional_name, prompt, needed_columns = formatter_template - if row_values[needed_columns[0]] not in (None, ""): + if row_values[needed_columns[0]] != "": prompt_values = {column: row_values[column] for column in needed_columns} formatter_values[optional_name] = prompt.format(**prompt_values) else: