Julenmendieta/milab 6494 support more embedding models

[julenmendieta]

Greptile Summary

This PR extends the Sequence Embeddings block from a single universal model (ESM-2) to a multi-model architecture, adding six specialist protein language models (CurrAb, AbLang2, VHHBERT, H3BERTa, TCR-BERT, PeptideCLM-2) with per-scope model selection via a new card-based UI. The data model migrates from a flat scope list with a global fidelity to per-card (scope × model) tasks (BlockDataV1 → V2), with a migration path that preserves ESM-2 for all existing projects.

• Model catalog (model/src/compat.ts): New single source of truth for model specs, scope↔model compatibility filtering (compatibleModels, compatibleScopes, isCompatible) and specialist-first recommendation logic (recommendedModel). Scopes gain isHeavy + receptor snapshots so the args lambda stays data-only.
• Workflow (models.lib.tengo, main.tpl.tengo, compute-embeddings.tpl.tengo): Central Tengo registry mirrors the TS catalog; per-task model routing uses resolveModelTag; isHeavy propagates to the Python step for AbLang2's heavy/light slot selection.
• Python step (main.py): Adds AbLang2Embedder, per-model input transforms (prepare_sequence), multi-family model loading (hf, hf-custom, ablang2), and per-model layer/pooling configuration controlled by --emb-layer and --pool-special-tokens.

Confidence Score: 3/5

The core multi-model architecture is well-designed, but CurrAb embeddings will be extracted from the wrong transformer layer (last vs. penultimate) until the registry discrepancy is resolved.

The CurrAb embLayer value in the Tengo registry (-1, last layer) directly contradicts the Python code's explicit statement that CurrAb uses penultimate (-2), and the workflow always overrides the Python default with the registry value. This means every CurrAb embedding run will silently use the wrong hidden layer. The rest of the architecture — the V1→V2 migration, scope compatibility logic, AbLang2 heavy/light slot routing, and the new card-based UI — is solid.

workflow/src/models.lib.tengo (CurrAb embLayer value) and ui/src/pages/EmbeddingCard.vue (sequence dropdown not filtered when model is chosen first)

Important Files Changed

Filename	Overview
model/src/compat.ts	New file: single source of truth for the embedding model catalog (EmbeddingModelSpec), scope↔model compatibility (compatibleModels, isCompatible, compatibleScopes), and specialist-first recommendation logic (recommendedModel). Logic is well-structured; SCEPTR is gated off safely via ENABLED_MODELS.
model/src/types.ts	EmbeddingModelId, ModelTag, EmbeddingCard, EmbeddingTask, BlockDataV1→V2 migration types, SelectedScope extended with isHeavy+receptor, ScopeConfig extended with receptor+paired. V1 type correctly frozen as SelectedScopeV1.
model/src/dataModel.ts	V1→V2 migration correctly seeds isHeavy:false + receptor:"IG" for ESM-2 (universal, ignores both fields) — safe default for existing projects. Init removes fidelity+selectedScopes in favour of empty embeddings list.
model/src/index.ts	Args projection validates each card's (scope, model) pair via isCompatible, deduplicates exact-duplicate tasks, and sorts by scope+model for deterministic args bytes. Error messages are user-friendly.
model/src/scopes.ts	Correctly adds isHeavy and receptor to all scope variants (peptide, CDR3, VDJRegion, Fv, scFv). deriveIsHeavy handles single-cell chain A/B and bulk IGHeavy. paired flag correctly drives VHH-vs-mAb recommendation default.
workflow/src/models.lib.tengo	Central registry for checkpoint attributes. CurrAb embLayer is -1 (last layer) but the Python docstring explicitly says CurrAb should use -2 (penultimate, like ESM-2). This inconsistency means CurrAb embeddings are extracted from the wrong layer.
workflow/src/main.tpl.tengo	Correctly iterates over embeddings tasks, resolves modelTag per task, passes isHeavy in the plan JSON for AbLang2 heavy/light slot routing, deduplicates columns in the stats counting pass, and builds distinct provenance traces per model.
workflow/src/compute-embeddings.tpl.tengo	Per-batch embedding body correctly pulls per-model recipe from the registry and has explicit asset-mount branches for all wired checkpoints. Correctly gates GPU memory request on device tier.
software/src_python/main.py	Adds AbLang2Embedder, prepare_sequence transform (CDR3 trim, spacing, SMILES), model-family routing, and per-model pool/layer config. Logic is sound; module docstring is ESM-2-specific and needs updating for multi-model context.
ui/src/pages/EmbeddingCard.vue	New card component with sequence→model filtering (model clears on incompatible scope pick). Model-first filtering of the sequence dropdown is absent; compatibleScopes is exported from compat.ts but not used here.
block/src/index.ts	New block facade file using ESM import.meta.url for the block-pack URL. BlockContract, BlockOutputs, BlockData aliases and block-named exports are all correctly defined.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    UI["EmbeddingCard UI\n(scope × model dropdowns)"] -->|writes EmbeddingCard| BD["BlockDataV2\n(embeddings: EmbeddingCard[])"]
    BD -->|args lambda| ARGS["BlockArgs\n(embeddings: EmbeddingTask[])"]
    ARGS -->|wf.body| WF["main.tpl.tengo\n(iterates tasks)"]
    WF -->|resolveModelTag| MLT["models.lib.tengo\ncentral registry"]
    MLT --> CE["compute-embeddings.tpl.tengo\n(per batch)"]
    CE -->|--model-family hf| EMB1["Embedder\n(AutoModel: ESM-2/CurrAb/VHHBERT/H3BERTa/TCR-BERT)"]
    CE -->|--model-family hf-custom| EMB2["Embedder trust_remote_code\n(PeptideCLM-2)"]
    CE -->|--model-family ablang2| EMB3["AbLang2Embedder\n(pip model + asset weights)"]
    EMB1 --> OUT["embeddings_{scope}.parquet\n(long-format: key x dim x value)"]
    EMB2 --> OUT
    EMB3 --> OUT
    COMPAT["model/src/compat.ts\nEMBEDDING_MODELS catalog"] -->|compatibleModels| UI
    COMPAT -->|isCompatible| ARGS
    SCOPES["model/src/scopes.ts\nbuildScopeConfig + isHeavy"] -->|ScopeConfig| UI

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    UI["EmbeddingCard UI\n(scope × model dropdowns)"] -->|writes EmbeddingCard| BD["BlockDataV2\n(embeddings: EmbeddingCard[])"]
    BD -->|args lambda| ARGS["BlockArgs\n(embeddings: EmbeddingTask[])"]
    ARGS -->|wf.body| WF["main.tpl.tengo\n(iterates tasks)"]
    WF -->|resolveModelTag| MLT["models.lib.tengo\ncentral registry"]
    MLT --> CE["compute-embeddings.tpl.tengo\n(per batch)"]
    CE -->|--model-family hf| EMB1["Embedder\n(AutoModel: ESM-2/CurrAb/VHHBERT/H3BERTa/TCR-BERT)"]
    CE -->|--model-family hf-custom| EMB2["Embedder trust_remote_code\n(PeptideCLM-2)"]
    CE -->|--model-family ablang2| EMB3["AbLang2Embedder\n(pip model + asset weights)"]
    EMB1 --> OUT["embeddings_{scope}.parquet\n(long-format: key x dim x value)"]
    EMB2 --> OUT
    EMB3 --> OUT
    COMPAT["model/src/compat.ts\nEMBEDDING_MODELS catalog"] -->|compatibleModels| UI
    COMPAT -->|isCompatible| ARGS
    SCOPES["model/src/scopes.ts\nbuildScopeConfig + isHeavy"] -->|ScopeConfig| UI

Loading

Prompt To Fix All With AI

Fix the following 3 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 3
workflow/src/models.lib.tengo:40
**CurrAb embLayer contradicts Python documentation**

The Tengo registry sets `embLayer: -1` (last layer) for CurrAb, but `main.py` explicitly states:

> "Per-model, following each model's OWN authoritative recipe: ESM-2 / **CurrAb use penultimate (-2, our validated default)**; VHHBERT / H3BERTa use the last layer (-1, their papers/code)"

Because `compute-embeddings.tpl.tengo` always passes `--emb-layer` explicitly from the registry (`string(m.embLayer)`), CurrAb embeddings are produced from the final transformer layer rather than the penultimate layer. ESM-2 and CurrAb are architecturally related (CurrAb is ESM-2-based), so both should extract from the same layer according to the Python doc. Using the wrong layer produces different — and likely lower-quality — representations. Either update the Tengo registry to `embLayer: -2`, or update the Python docstring and validate that `-1` is intentional for CurrAb.

### Issue 2 of 3
ui/src/pages/EmbeddingCard.vue:35-37
**Sequence dropdown not filtered when a model is selected first**

`sequenceOptions` maps all `config.options` regardless of any selected model, so model-first selection is not filtered. The `compat.ts` module docstring describes "bidirectional filtering" and `compatibleScopes` is exported precisely for this purpose, but it is never called here.

Concretely: a user who selects VHHBERT or H3BERTa (heavy-chain-only specialists) will still see light-chain CDR3 and VDJRegion scopes in the sequence dropdown. Choosing one silently clears the model (via the `keep` guard in `onSequence`), which is confusing. Filtering `sequenceOptions` with `compatibleScopes` when `card.model` is set would prevent incompatible pairs from ever being shown.

### Issue 3 of 3
software/src_python/main.py:1-15
**Module docstring is ESM-2-specific and doesn't reflect multi-model support**

The top-level docstring still describes the script as loading "ONE protein language model (the exact checkpoint the workflow chose for the runtime device tier…)" and specifically references "ESM-2 650M (fp16), ESM-2 150M (fp32)" and "HuggingFace ESM-2-style checkpoint". Now that the script supports CurrAb, VHHBERT, H3BERTa, TCR-BERT, PeptideCLM-2 (HF), and AbLang2 (custom pip), the docstring is misleading. It also does not mention the new `--model-family`, `--pair-mode`, `--emb-layer`, or `--pool-special-tokens` arguments that control per-model behavior.

_{Reviews (1): Last reviewed commit: "Changeset" | Re-trigger Greptile}

Greptile also left 2 inline comments on this PR.

Context used:

• Context used - Terms is a types in codebase. Provide the list of ... (source)

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.