From 64020c321d7c1c8d7683db995cf5a9e74fd845da Mon Sep 17 00:00:00 2001 From: Tiara Rodney Date: Wed, 17 Jun 2026 14:03:52 +0200 Subject: [PATCH] test: annotate the sft test helpers --- tests/unit/test_sft.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_sft.py b/tests/unit/test_sft.py index f2e84b4..d24eef0 100644 --- a/tests/unit/test_sft.py +++ b/tests/unit/test_sft.py @@ -15,7 +15,7 @@ class FakeTok: """Additive chat template: each turn renders to `` tokens... ``; the generation prompt appends ````.""" - 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