Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 75 additions & 0 deletions tests/python/test_to_sharegpt_optional_none.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
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!"]
9 changes: 8 additions & 1 deletion unsloth/chat_templates.py
Original file line number Diff line number Diff line change
Expand Up @@ -2163,7 +2163,14 @@ def __combined_prompt_processor__(examples):

_, optional_name, prompt, needed_columns = formatter_template
if row_values[needed_columns[0]] not in (None, ""):
prompt_values = {column: row_values[column] for column in needed_columns}
# Coerce missing (None) columns to "" so they do not render as the
# literal string "None" in the emitted text. A [[...]] block may
# reference several columns; only the first gates the block, so a
# later column can still be None here.
prompt_values = {
column: ("" if row_values[column] is None else row_values[column])
for column in needed_columns
}
formatter_values[optional_name] = prompt.format(**prompt_values)
else:
formatter_values[optional_name] = ""
Expand Down
Loading