diff --git a/flexeval/core/language_model/hf_lm.py b/flexeval/core/language_model/hf_lm.py index 46c7d636..0d2a4770 100644 --- a/flexeval/core/language_model/hf_lm.py +++ b/flexeval/core/language_model/hf_lm.py @@ -383,15 +383,36 @@ def _batch_complete_text( # We strip the input text and stop sequences from the output text. lm_outputs: list[LMOutput] = [] - for generated_tensor in generated_tokens: + for attention_mask, generated_tensor in zip(model_inputs["attention_mask"], generated_tokens): input_tensor = generated_tensor[:input_token_length] output_tensor = generated_tensor[input_token_length:] - input_tokens = [t for t in input_tensor.tolist() if t != self.tokenizer.pad_token_id] - output_tokens = [t for t in output_tensor.tolist() if t != self.tokenizer.pad_token_id] + # Use the attention mask (rather than filtering by `pad_token_id`) to strip the + # left-padding from the input, since `pad_token_id` may equal `eos_token_id` (the + # tokenizer has no native pad token), in which case a value-based filter would also + # strip genuine EOS tokens that are part of the actual prompt content. + input_tokens = [t for t, m in zip(input_tensor.tolist(), attention_mask.tolist()) if m != 0] + + # Find where generation actually stopped: the first occurrence of a stop token id in + # the output. Everything from that point on (the stop token itself and the padding + # that follows it) is discarded. Note that `pad_token_id` may equal one of the + # `stop_token_ids` (when the tokenizer has no native pad token), so a value-based pad + # filter would incorrectly strip the genuine stop token as well, making it impossible + # to tell "stop" from "length" below. + stop_token_pos = next( + (i for i, t in enumerate(output_tensor.tolist()) if t in stop_token_ids), + None, + ) + if stop_token_pos is not None: + finish_reason = "stop" + output_tokens = output_tensor[:stop_token_pos].tolist() + else: + finish_reason = "length" + # Fall back to filtering by `pad_token_id`, as a safety net for models with a + # dedicated pad token (distinct from eos) that may still pad the output. + output_tokens = [t for t in output_tensor.tolist() if t != self.tokenizer.pad_token_id] decoded_text = decode_for_lm_continuation(output_tokens, input_tokens, self.tokenizer) - finish_reason = "length" for stop_seq in stop_sequences: idx = decoded_text.find(stop_seq) if idx != -1: diff --git a/tests/core/language_model/test_hf_lm.py b/tests/core/language_model/test_hf_lm.py index 1ed6547c..e33b8958 100644 --- a/tests/core/language_model/test_hf_lm.py +++ b/tests/core/language_model/test_hf_lm.py @@ -88,6 +88,31 @@ def chat_lm_without_system_message(model_name: str = "sbintuitions/tiny-lm-chat" ) +class _StubGenerateModel: + """A fake `model` object that stands in for the lazily-loaded HF model. + + `HuggingFaceLM._batch_complete_text` calls `self.model.generate(**model_inputs, **gen_kwargs)`. + By assigning an instance of this class directly to `lm.model` before calling `complete_text`, + the `@load_model` decorator's lazy-loading branch is skipped (since `self.model is not None`), + and `generate()` deterministically returns the real (padded) input ids concatenated with a + pre-specified sequence of "generated" token ids per batch row. This lets us test the + post-processing logic in `_batch_complete_text` without depending on non-deterministic + real generation. + """ + + def __init__(self, generated_tokens: list[list[int]]) -> None: + self.device = torch.device("cpu") + self._generated_tokens = generated_tokens + + def generate(self, **kwargs) -> torch.Tensor: + input_ids = kwargs["input_ids"] + rows = [] + for i, ids in enumerate(input_ids): + gen = torch.tensor(self._generated_tokens[i], dtype=ids.dtype) + rows.append(torch.cat([ids, gen])) + return torch.stack(rows) + + class TestHuggingFaceLM(BaseLanguageModelTest): @pytest.fixture def lm(self, lm: HuggingFaceLM) -> LanguageModel: @@ -307,6 +332,119 @@ def test_if_stop_sequences_work_as_expected(chat_lm: HuggingFaceLM) -> None: assert response.finish_reason == "length" +def test_finish_reason_is_stop_when_eos_generated_with_padless_tokenizer() -> None: + """Regression test for the bug where `finish_reason` was always "length" for tokenizers + without a native pad_token (e.g. GPT-2/Llama family), because `tokenize_text_for_lm_prefix` + falls back to `pad_token = eos_token`, and the output post-processing used to strip *all* + tokens equal to `pad_token_id` -- including the genuine EOS token that caused generation to + stop -- before the text-level stop-sequence search ran. + """ + lm = HuggingFaceLM( + model="tokyotech-llm/Swallow-7b-instruct-hf", + tokenizer_kwargs={"use_fast": False}, + default_gen_kwargs={"do_sample": False}, + model_limit_tokens=None, + ) + assert lm.tokenizer.pad_token is None # sanity check: this tokenizer has no native pad_token + + eos_id = lm.tokenizer.eos_token_id + # Simulate: 2 generated content tokens, then EOS (the real stop), then EOS-valued padding. + generated_tokens = [100, 200, eos_id, eos_id, eos_id] + lm.model = _StubGenerateModel(generated_tokens=[generated_tokens]) + + output = lm.complete_text("hello", max_new_tokens=5) + assert output.finish_reason == "stop" + + +def test_finish_reason_is_length_when_max_new_tokens_reached_without_eos() -> None: + """Sanity check: without any EOS token among the generated tokens, `finish_reason` should + remain "length" both before and after the fix.""" + lm = HuggingFaceLM( + model="tokyotech-llm/Swallow-7b-instruct-hf", + tokenizer_kwargs={"use_fast": False}, + default_gen_kwargs={"do_sample": False}, + model_limit_tokens=None, + ) + assert lm.tokenizer.pad_token is None + + generated_tokens = [100, 200, 300] # no EOS among these + lm.model = _StubGenerateModel(generated_tokens=[generated_tokens]) + + output = lm.complete_text("hello", max_new_tokens=3) + assert output.finish_reason == "length" + + +def test_finish_reason_mixed_batch_with_padless_tokenizer() -> None: + """A batch with one sequence that stops early (EOS + EOS-padding) and one that runs to + max_new_tokens should independently get "stop" and "length" as their finish_reason.""" + lm = HuggingFaceLM( + model="tokyotech-llm/Swallow-7b-instruct-hf", + tokenizer_kwargs={"use_fast": False}, + default_gen_kwargs={"do_sample": False}, + model_limit_tokens=None, + ) + assert lm.tokenizer.pad_token is None + + eos_id = lm.tokenizer.eos_token_id + stopped_row = [100, 200, eos_id, eos_id] + length_row = [100, 200, 300, 400] + lm.model = _StubGenerateModel(generated_tokens=[stopped_row, length_row]) + + outputs = lm.complete_text(["hello", "world"], max_new_tokens=4) + assert outputs[0].finish_reason == "stop" + assert outputs[1].finish_reason == "length" + + +def test_input_side_pad_filter_preserves_genuine_eos_in_prompt() -> None: + """Regression test for the input-side counterpart of the same bug: filtering the prompt's + tokens by `token != pad_token_id` removes genuine EOS tokens that are part of the actual + prompt content (not just left-padding), corrupting `decode_for_lm_continuation`'s boundary + computation. The fix filters using `attention_mask` instead, which correctly identifies only + the left-padding positions. + + We use a SentencePiece-based tokenizer (no native pad_token) because the corruption manifests + as a lost leading space: SentencePiece's decode() strips a leading space only from the very + first token of the tokens passed to it, so if the buggy filter empties out the "prompt" token + list entirely, the leading space that should have belonged to the prompt gets attributed to + the continuation and stripped away instead. + """ + tokenizer_name = "tokyotech-llm/Swallow-7b-instruct-hf" + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=False) + assert tokenizer.pad_token is None # sanity check: this tokenizer has no native pad_token + + # item 0: the entire prompt content is just the eos_token itself (so after left-padding with + # pad_token==eos_token, the row is *entirely* composed of eos-valued tokens, both the fake + # padding and the one genuine content token). + # item 1: a long, unrelated prompt to force item 0 to be left-padded. + text_list = [ + tokenizer.eos_token, + "これはとても長い日本語の文章です、パディングを発生させるために十分な長さがあります。", + ] + + # Compute the padded input ids exactly as `_batch_complete_text` will internally, so we know + # what the stub model will receive and can compute the expected (correct) continuation text. + model_inputs = tokenize_text_for_lm_prefix(text_list, tokenizer, add_special_tokens=False) + row0_ids = model_inputs.input_ids[0].tolist() + row0_mask = model_inputs.attention_mask[0].tolist() + correct_input_tokens_row0 = [t for t, m in zip(row0_ids, row0_mask) if m == 1] + assert correct_input_tokens_row0, "the genuine eos token in the prompt should be preserved" + + generated_tokens = tokenizer("です。", add_special_tokens=False)["input_ids"] + expected_text_row0 = decode_for_lm_continuation(generated_tokens, correct_input_tokens_row0, tokenizer) + + lm = HuggingFaceLM( + model=tokenizer_name, + tokenizer_kwargs={"use_fast": False}, + default_gen_kwargs={"do_sample": False}, + model_limit_tokens=None, + ) + lm.model = _StubGenerateModel(generated_tokens=[generated_tokens, generated_tokens]) + + outputs = lm.complete_text(text_list, max_new_tokens=len(generated_tokens)) + assert outputs[0].text == expected_text_row0 + assert outputs[0].text.startswith(" ") # the leading space that belongs to the prompt/continuation boundary + + def test_if_gen_kwargs_work_as_expected() -> None: lm = HuggingFaceLM(model="sbintuitions/tiny-lm", default_gen_kwargs={"max_new_tokens": 1}) # check if the default gen_kwargs is used and the max_new_tokens is 1