diff --git a/tests/unit/model_bridge/test_generate_fallback_flags.py b/tests/unit/model_bridge/test_generate_fallback_flags.py index 46d5ac94f..22df352bb 100644 --- a/tests/unit/model_bridge/test_generate_fallback_flags.py +++ b/tests/unit/model_bridge/test_generate_fallback_flags.py @@ -1,7 +1,8 @@ -"""generate()'s stateful hf_generate fallback must forward its stop flags.""" +"""Stateful hf_generate fallback preserves stop flags and sampling penalties.""" from types import SimpleNamespace +import pytest import torch from transformers import GPT2Config, GPT2LMHeadModel @@ -108,3 +109,59 @@ def spy(input_ids, **kwargs): ) # HF takes int or list; any other sequence must arrive normalized. assert captured["eos_token_id"] == [5, 7] + + +@pytest.mark.parametrize("penalty", [0.8, 1.0, 1.5]) +def test_stateful_fallback_forwards_repetition_penalty(monkeypatch, penalty): + bridge = _tiny_bridge() + bridge.cfg.is_stateful = True + bridge.tokenizer = SimpleNamespace(eos_token_id=0) + captured: dict = {} + + def spy(input_ids, **kwargs): + captured.update(kwargs) + return input_ids + + monkeypatch.setattr(bridge.original_model, "generate", spy) + bridge.generate( + torch.tensor([[1, 2, 3]]), + use_past_kv_cache=False, + repetition_penalty=penalty, + do_sample=False, + verbose=False, + ) + assert captured["repetition_penalty"] == penalty + + +@pytest.mark.parametrize("penalty", [-0.5, 0.5]) +def test_stateful_fallback_rejects_frequency_penalty(monkeypatch, penalty): + bridge = _tiny_bridge() + bridge.cfg.is_stateful = True + captured = _capture_hf_generate(bridge, monkeypatch) + with pytest.raises(NotImplementedError, match="freq_penalty.*use_past_kv_cache=True"): + bridge.generate( + torch.tensor([[1, 2, 3]]), + use_past_kv_cache=False, + freq_penalty=penalty, + verbose=False, + ) + assert not captured + + +def test_stateful_fallback_repetition_penalty_matches_hf_generation(): + bridge = _tiny_bridge() + bridge.cfg.is_stateful = True + bridge.tokenizer = SimpleNamespace(eos_token_id=0) + tokens = torch.tensor([[1, 2, 3]]) + kwargs = dict(max_new_tokens=4, do_sample=False, stop_at_eos=False) + expected = bridge.hf_generate(tokens, repetition_penalty=2.0, **kwargs) + unpenalized = bridge.hf_generate(tokens, repetition_penalty=1.0, **kwargs) + actual = bridge.generate( + tokens, + repetition_penalty=2.0, + use_past_kv_cache=False, + verbose=False, + **kwargs, + ) + torch.testing.assert_close(actual, expected) + assert not torch.equal(actual, unpenalized) diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 4a4ede408..f4c0c43d1 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -2998,7 +2998,9 @@ def generate( top_k: Number of tokens to sample from. If None, sample from all tokens top_p: Probability mass to sample from. If 1.0, sample from all tokens temperature: Temperature for sampling. Higher values will make the model more random - freq_penalty: Frequency penalty for sampling - how much to penalise previous tokens + freq_penalty: Frequency penalty for sampling - how much to penalise previous tokens. + Nonzero values require the hooked generation loop; the stateful/SSM + HF fallback with use_past_kv_cache=False raises NotImplementedError. repetition_penalty: HuggingFace-style repetition penalty. Values > 1.0 discourage repetition by dividing positive logits and multiplying negative logits for previously seen tokens. Default 1.0 (no penalty). @@ -3367,10 +3369,16 @@ def generate( ) if is_stateful_model and not use_stateful_cache: + if freq_penalty != 0.0: + raise NotImplementedError( + "freq_penalty requires use_past_kv_cache=True for stateful/SSM models. " + "The HF generation fallback has no equivalent frequency penalty." + ) hf_kwargs: dict[str, Any] = { "max_new_tokens": max_new_tokens, "do_sample": do_sample, "temperature": temperature, + "repetition_penalty": repetition_penalty, # hf_generate defaults this to True, so omitting it would make # generate(stop_at_eos=False) silently stop at EOS on this path. "stop_at_eos": stop_at_eos,