[jahnzh]

Summary

Resolves #23.

This PR implements the fixes for ESM-C logits extraction and repairs the codebase-wide logits extraction pipeline.

Changes Made

1. ESMCEmbedder._compute_outputs:

   • Implemented _compute_outputs in src/pepe/embedders/huggingface_embedder.py to forward passes through AutoModelForMaskedLM and return LM-head logits on CPU with (batch, seq, vocab) shape.
   • Initialized self.vocab_size = 64 (reflecting Transformer Engine / FP8 64-dimensional padded vocabulary).

2. Removal of Erroneous .permute(2, 0, 1):

   • Removed .permute(2, 0, 1) from ESM2Embedder._compute_outputs and ESMEmbedder._compute_outputs. Logits now preserve standard (batch, max_input_length, vocab_size) layout across all embedders.

3. BaseEmbedder Codebase-Wide Fixes:

   • _set_output_objects: Added helper _get_vocab_size() to dynamically detect vocab dimension (e.g. 64 for ESM-C, 33 for ESM-2/1). Configured self.logits["shape"] as (num_sequences, max_input_length, vocab_size) (or flattened (num_sequences, max_input_length * vocab_size)).
   • _extract_logits: Fixed extraction to operate on 3D tensors (batch, seq, vocab). Eliminated layer - 1 indexing bug that previously mistook vocabulary indices for layer IDs.
   • _safe_compute OOM recovery: Verified dim=0 batch-halving concatenation works as intended with unpermuted tensors.
   • export_to_disk: Updated flatten logic to support flattening logits alongside per-token embeddings when --flatten is specified.

4. Model Inspector (model_selecter.py):

   • Updated report_model() so inspect_model / pepe --check_model reports Logits: yes for ESMCEmbedder.

5. Tests & Documentation:

   • Added unit tests in src/tests/test_esmc.py validating report_model reports logits availability and ESMCEmbedder._compute_outputs returns properly shaped (batch, seq, 64) logits.
   • Added test_logits to src/tests/test_esmc_modes.py asserting output ndim and 64-dim vocab.
   • Updated regression tests in src/tests/test_streaming_roundtrip.py validating exact match between streaming memmapped disk arrays and in-memory arrays for 3D logits (batch, seq, 33).
   • Updated CLAUDE.md architecture reference.

Scientific Caveat & Note on ESM-C Logits

• ESM-C was trained with 100% [MASK] replacement during masked language modeling.
• Single-pass unmasked logits are made available as raw model outputs, but users performing zero-shot variant effect prediction or CDR3 sequence recovery should be aware that unmasked logits exhibit collapse / degenerate behavior (~35–67% recovery) compared to ESM-2. Masked marginal scoring remains recommended for variant effect prediction.

[Copilot]

🟡 Changes recommended

Unresolved logits reconstruction, custom vocabulary sizing, and OOM retry issues can corrupt or prevent logits extraction.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Fixes ESM-C logits extraction and standardizes logits handling across embedders as 3D batch/sequence/vocabulary tensors.

Changes:

• Adds ESM-C logits computation and model-inspector support.
• Repairs logits extraction, sizing, flattening, and streaming preallocation.
• Adds regression tests and documents ESM-C scoring caveats.

File summaries

File	Summary
src/tests/test_streaming_roundtrip.py	Validates 3D logits streaming round trips.
src/tests/test_esmc.py	Tests ESM-C logits and model reporting.
src/tests/test_esmc_modes.py	Tests ESM-C logits output dimensions.
src/pepe/model_selecter.py	Reports ESM-C logits availability.
src/pepe/embedders/huggingface_embedder.py	Implements ESM-C logits computation.
src/pepe/embedders/esm_embedder.py	Standardizes logits layout.
src/pepe/embedders/custom_embedder.py	Adds custom vocabulary sizing metadata.
src/pepe/embedders/base_embedder.py	Updates logits extraction, reconstruction, and output sizing.
CLAUDE.md	Documents logits behavior and ESM-C caveats.

Review details

Suppressed comments (4)

CLAUDE.md:71

• The caveat is documented only in CLAUDE.md; the user-facing README still calls logits merely experimental and the CLI help has no warning. Because unmasked ESM-C logits can produce scientifically invalid variant scores, the required warning will not reach users invoking pepe or the public API. Add the masking/scoring caveat and 64-wide vocabulary note to the user documentation and CLI help.

`_set_output_objects()` defines one entry per output type — `per_token`, `mean_pooled`, `substring_pooled`, `attention_head`, `attention_layer`, `attention_model`, `logits` — each a dict of `{output_data, method (_extract_*), output_dir, shape}`. `extract_embeddings` picks which run. Adding an output type means adding both the dict entry and the matching `_extract_*` method. Files are saved as `.npy` under `<output_path>/<model_name>/<output_type>/`, plus a `*_idx.csv` mapping row index → sequence id. Note `logits` is produced by ESM-2 and ESM-C (via `AutoModelForMaskedLM`); the generic HuggingFace path drops it. (Note: ESM-C outputs 64-dimensional unmasked single-pass logits with indices 4–28 supervised; zero-shot variant effect prediction scientifically requires masked marginal scoring).

src/pepe/embedders/base_embedder.py:180

• The new 3-D preallocation/write path is covered by the ESM-2 round-trip test, but the ESM-C tests only cover in-memory output and a mocked forward pass; the real ESM-C integration is skipped unless ESMC_TEST=1. Since ESM-C uses a distinct 64-wide model and this PR changes shared streaming memmaps, add an ESM-C streaming round-trip or focused preallocation/write test to catch shape mismatches.

                (
                    self.num_sequences,
                    self.max_input_length,
                    self.vocab_size,
                )
                if not self.flatten
                else (
                    self.num_sequences,
                    self.max_input_length * self.vocab_size,
                )

src/pepe/embedders/huggingface_embedder.py:655

• This returns representations as a layer-keyed dict, but BaseEmbedder._safe_compute() concatenates recursive o[1] values with torch.cat as though they were tensors. If logits are requested together with embeddings and an OOM retry occurs, this new ESM-C path raises a TypeError instead of recovering. Update the retry path to concatenate each representation layer independently.

        if return_embeddings:
            representations = {
                layer: outputs.hidden_states[layer]
                .to(self._precision_to_dtype(self.precision, "torch"))
                .cpu()
                for layer in self.layers

src/tests/test_streaming_roundtrip.py:155

• These assertions cover the ESM-2 33-wide path only. The PR also introduces ESM-C's 64-wide logits and changes memmap preallocation to depend on that dimension, but the ESM-C tests are in-memory and do not exercise streaming writes or round-trip shape equality. Add an opt-in ESM-C streaming/in-memory regression so the new padded-vocabulary path is tested.

        if output_type == "logits":
            assert reference.ndim == 3, f"Expected 3D logits, got {reference.ndim}"
            assert reference.shape[-1] == 33, f"Expected ESM-2 vocab size 33, got {reference.shape[-1]}"
            assert streaming.shape == reference.shape

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

---

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

[jahnzh]

Summary of Updates Addressing Review Comments & CI

1. ESM-C Documentation & CLI Warnings:

   • Added documentation in README.md and CLI help in src/pepe/parse_arguments.py detailing ESM-C's 64-dimensional padded output space (FP8 / Transformer Engine layout, supervised tokens at indices 4–28) and noting that single-pass unmasked logits exhibit collapse / degenerate behavior compared to ESM-2, with masked marginal scoring recommended for zero-shot variant effect prediction.

2. OOM Retry Recovery for Layer-Keyed Dictionaries:

   • Updated BaseEmbedder._safe_compute() to concatenate representations per layer across recursed chunks when representations is returned as a dictionary, preventing TypeError during OOM recovery.

3. Custom Vocabulary Sizing:

   • Configured CustomModelWrapper in src/pepe/embedders/custom_embedder.py to read config.get('vocab_size', 21) instead of hardcoding 21.
   • Updated CustomEmbedder._initialize_model() to derive self.vocab_size from model.vocab_size / config / tokenizer with 21 as fallback, ensuring disk preallocation matches output shapes.

4. Chunk Reconstruction with discard_padding=True and flatten=True:

   • Updated BaseEmbedder._reconstruct_chunks() to support discard_padding=True, slicing chunks from index 0 without skipping the first residue of subsequent chunks when special tokens have already been stripped.
   • Deferred flattening in _extract_logits() and _extract_per_token() when sequence chunking is active, flattening reconstructed sequences after stitching so 2D token dimensions are preserved during reconstruction.

5. Type Checker CI Fixes:

   • Resolved all type annotations and indexability errors in BaseEmbedder._safe_compute() and CustomModelWrapper.
   • Verified that all GitHub Actions workflows (Pre-commit linting, Unit tests on Python 3.10 and 3.11, Mypy type checker, and Integration tests) pass cleanly.