tiararodney/sekft
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

test: annotate the sft test helpers

Browse source
This commit is contained in:
Tiara Rodney 2026-06-17 14:03:52 +02:00
parent 9397280e6f
commit 64020c321d
Signed by: tiara
GPG key ID: 5CD8EC1D46106723
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
tests/unit/test_sft.py
View file

@ -15,7 +15,7 @@ class FakeTok:
"""Additive chat template: each turn renders to ``<role> tokens... </e>``;
the generation prompt appends ``<assistant>``."""
def apply_chat_template(self, msgs: list[dict], add_generation_prompt: bool = False,
def apply_chat_template(self, msgs: list[dict[str, str]], add_generation_prompt: bool = False,
return_tensors: Any = None) -> list[str]:
toks: list[str] = []
for m in msgs:
@ -65,7 +65,7 @@ def test_mask_trains_assistant_turns_only() -> None:
def test_mask_raises_on_non_additive_template() -> None:
class BadTok:
def apply_chat_template(self, msgs: list[dict], add_generation_prompt: bool = False,
def apply_chat_template(self, msgs: list[dict[str, str]], add_generation_prompt: bool = False,
return_tensors: Any = None) -> list[int]:
return list(range(len(msgs), 0, -1)) # reversed: prefixes do not nest