Keep the generation regression harness outside the model repository

tests/test_generation_positions.py

@@ -1,93 +0,0 @@
-"""Run with pytest under Transformers 4.57.6 and 5.22.1; no model weights needed."""
-
-import importlib.util
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-
-ROOT = Path(__file__).resolve().parents[1]
-SPEC = importlib.util.spec_from_file_location(
-    "ripple_generation_regression", ROOT / "modeling_ricursive.py",
-    submodule_search_locations=[str(ROOT)],
-)
-MODEL = importlib.util.module_from_spec(SPEC)
-sys.modules[SPEC.name] = MODEL
-SPEC.loader.exec_module(MODEL)
-
-
-@pytest.fixture
-def model():
-    torch.manual_seed(19)
-    config = MODEL.RicursiveConfig(
-        vocab_size=97, hidden_size=32, intermediate_size=64,
-        num_hidden_layers=5, num_attention_heads=4, num_key_value_heads=2,
-        head_dim=8, sliding_window=16, max_position_embeddings=128,
-        pad_token_id=0, eos_token_id=None, use_valres=True,
-    )
-    return MODEL.RicursiveForCausalLM(config).eval()
-
-
-@pytest.mark.parametrize("use_cache", [True, False])
-@pytest.mark.parametrize("custom_positions", [False, True])
-@torch.inference_mode()
-def test_explicit_positions_extend_during_generation(model, use_cache, custom_positions):
-    prompt = torch.tensor([[3, 4, 5, 6, 7, 8], [0, 0, 3, 4, 5, 6]])
-    mask = prompt.ne(0).long()
-    positions = (mask.cumsum(-1) - 1).masked_fill(mask == 0, 0)
-    if custom_positions:
-        positions = torch.tensor([[2, 4, 5, 8, 9, 12], [0, 0, 17, 19, 20, 24]])
-    original_positions = positions.clone()
-    expected = prompt.clone()
-    expected_mask, expected_positions = mask.clone(), positions.clone()
-    # Independent full-forward decoding checks both offsets and newly added positions.
-    for _ in range(4):
-        logits = model(expected, attention_mask=expected_mask,
-                       position_ids=expected_positions, use_cache=False).logits
-        expected = torch.cat([expected, logits[:, -1].argmax(-1, keepdim=True)], dim=-1)
-        expected_mask = torch.cat([expected_mask, torch.ones((2, 1), dtype=mask.dtype)], dim=-1)
-        expected_positions = torch.cat([expected_positions, expected_positions[:, -1:] + 1], dim=-1)
-
-    observed = []
-    def capture(_module, _args, kwargs):
-        observed.append(kwargs["position_ids"].clone())
-    handle = model.register_forward_pre_hook(capture, with_kwargs=True)
-    try:
-        actual = model.generate(prompt, attention_mask=mask, position_ids=positions,
-                                max_new_tokens=4, do_sample=False, use_cache=use_cache)
-    finally:
-        handle.remove()
-    torch.testing.assert_close(actual, expected)
-    torch.testing.assert_close(positions, original_positions)
-    torch.testing.assert_close(observed[0], original_positions)
-    for step, seen in enumerate(observed[1:], start=1):
-        future = original_positions[:, -1:] + torch.arange(1, step + 1)
-        wanted = future[:, -1:] if use_cache else torch.cat([original_positions, future], dim=-1)
-        torch.testing.assert_close(seen, wanted)
-
-
-@torch.inference_mode()
-def test_explicit_positions_with_beam_search(model):
-    prompt = torch.tensor([[3, 4, 5, 6, 7, 8]])
-    expected = model.generate(prompt, num_beams=3, max_new_tokens=4, do_sample=False)
-    actual = model.generate(prompt, position_ids=torch.arange(6).unsqueeze(0),
-                            num_beams=3, max_new_tokens=4, do_sample=False)
-    torch.testing.assert_close(actual, expected)
-
-
-@pytest.mark.parametrize("num_new_tokens", [1, 3])
-@torch.inference_mode()
-def test_position_update_appends_exactly_the_accepted_token_count(model, num_new_tokens):
-    prompt = torch.tensor([[3, 4, 5, 6]])
-    positions = torch.tensor([[7, 8, 10, 13]])
-    output = model(prompt, position_ids=positions, use_cache=True)
-    updated = model._update_model_kwargs_for_generation(
-        output,
-        {"position_ids": positions, "attention_mask": torch.ones_like(prompt),
-         "cache_position": torch.arange(4), "use_cache": True},
-        num_new_tokens=num_new_tokens,
-    )
-    expected = torch.cat([positions, positions[:, -1:] + torch.arange(1, num_new_tokens + 1)], dim=-1)
-    torch.testing.assert_close(updated["position_ids"], expected)