[jahnzh]

Summary

This PR adds an extensive regression and contract test suite to prevent dimension mismatch and pipeline errors across all output extraction methods and execution modes.

Key Changes

1. Contract Test Suite (src/tests/test_output_methods_contract.py):

   • Shape & Preallocation Invariant: Verifies disk preallocation registered in _set_output_objects() exactly matches written arrays and in-memory shapes for all output types (mean_pooled, per_token, logits, attention_head, attention_layer, attention_model).
   • Flatten Contract: Ensures flatten=True yields valid 2D matrices across all output types in both streaming and in-memory modes.
   • Multi-layer Indexing Integrity: Validates simultaneous extraction of multiple layers without non-layer slicing bugs.
   • Padding Discard Contract: Ensures variable-length sequences with discard_padding=True cleanly strip padding tokens across token-level and logit outputs.
   • Splitting & Chunk Reconstruction: Validates that split_long_sequences=True properly stitches multi-chunk sequences across per_token, logits, and mean_pooled.
   • Substring Pooling: Confirms that substring_pooled correctly preserves (N, D) representations.
   • Vocab Size Resolver: Verifies _get_vocab_size() safely resolves vocab sizes across model backends, alphabets, and configs.

2. BaseEmbedder Hardening (src/pepe/embedders/base_embedder.py):

   • Updated _prepare_tensor to safely handle variable-length sequence arrays when saving unpadded/chunked embeddings without raising RuntimeError: stack expects each tensor to be equal size.

Test Coverage

• Passes pytest src/tests/test_output_methods_contract.py (7/7 tests pass).
• Passes ruff check, ruff format --check, and mypy src/ with 0 issues.

[Copilot]

🟡 Changes recommended

A critical export error-handling issue and a moderate contract-test gap remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds contract tests for output extraction behavior and hardens variable-length tensor export handling.

Changes:

• Added regression tests for shapes, flattening, padding, splitting, pooling, and vocabulary resolution.
• Added variable-length tensor serialization fallback.

File summaries

File	Summary	Review findings
src/tests/test_output_methods_contract.py	Adds output-method contract and regression tests.	Moderate (1 vote): strengthen reconstructed mean-value assertions. Nits (1, 1, and 3 votes): cover in-memory attention-head shapes, all in-memory flattened outputs, and config-only vocabulary resolution.
src/pepe/embedders/base_embedder.py	Handles variable-length tensors during export.	Critical (1 vote; lines 936–937): avoid catching every RuntimeError from torch.stack; distinguish unequal-size failures from other tensor errors.

Review details

Suppressed comments (4)

src/pepe/embedders/base_embedder.py:939

• When reconstructed sequences have different lengths, this fallback returns a dtype=object array, and export_to_disk() then writes it with np.save as a pickle-backed .npy. Normal consumers (including the repository's np.load calls, which do not set allow_pickle=True) cannot load that file and raise ValueError: Object arrays cannot be loaded when allow_pickle=False. Please pad reconstructed outputs to a common numeric shape or define and use a supported non-pickle variable-length format instead of silently producing an unusable disk artifact.

            if flatten:
                return np.array([t.numpy().flatten() for t in data_list], dtype=object)
            return np.array([t.numpy() for t in data_list], dtype=object)

src/tests/test_output_methods_contract.py:130

• Although expected_shapes includes attention_head, the in-memory assertions here cover every listed output except per-head attention; only the streamed head files are checked. A regression that changes res_mem["attention_head"][LAYER][head] can therefore pass this contract test. Add a shape assertion for each head in the in-memory result.

            attn_layer_arr = np.stack(
                [t.numpy() for t in res_mem["attention_layer"][LAYER]]
            )
            self.assertEqual(attn_layer_arr.shape, expected_shapes["attention_layer"])

src/tests/test_output_methods_contract.py:208

• The flattened contract is advertised for all output types, but the in-memory portion only checks logits and per_token. Mean-pooled and all three attention outputs are validated only on disk, so an in-memory flattening regression in those paths would still pass. Add corresponding in-memory shape checks for the remaining output types.

            # In-memory flattened checks
            self.assertEqual(
                np.stack([t.numpy() for t in res_mem["logits"][LAYER]]).shape,
                expected_flattened_shapes["logits"],
            )
            self.assertEqual(
                np.stack([t.numpy() for t in res_mem["per_token"][LAYER]]).shape,
                expected_flattened_shapes["per_token"],
            )

src/tests/test_output_methods_contract.py:353

• These assertions only check the reconstructed mean vector's rank, not its values. In the no-discard path, _extract_mean_pooled excludes special tokens, while _reconstruct_chunks recomputes full_per_token.mean(0), so BOS/EOS can be included without this test failing. Assert that the reconstructed mean equals the residue-only mean of per_token[1:-1] (and similarly cover the discard case).

        self.assertEqual(res_no_discard["per_token"][LAYER][0].shape, (32, EMBED_DIM))
        self.assertEqual(res_no_discard["logits"][LAYER][0].shape, (32, VOCAB_SIZE))

• Files reviewed: 2/2 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.