Skip to content

Expose BERT token-type embeddings through TransformerBridge - #1664

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:dev-4.xfrom
LarryHu0217:codex/bert-token-type-embed-1662
Aug 13, 2026
Merged

Expose BERT token-type embeddings through TransformerBridge#1664
jlarson4 merged 1 commit into
TransformerLensOrg:dev-4.xfrom
LarryHu0217:codex/bert-token-type-embed-1662

Conversation

@LarryHu0217

Copy link
Copy Markdown

Description

Expose BERT token-type (segment) embeddings as the bridge-native token_type_embed.hook_out hook.

The BERT adapter now maps bert.embeddings.token_type_embeddings through an EmbeddingBridge. The regression test boots bert-base-cased, captures the new hook with run_with_cache, and compares it directly with Hugging Face's token-type embedding output.

Fixes #1662

Type of change

  • New feature (non-breaking change which adds functionality)

Validation

  • uv run pytest tests/unit/model_bridge/supported_architectures/test_bert_adapter.py -q --tb=short -o addopts=: 22 passed
  • make unit-test: 5,260 passed, 56 skipped, 54 deselected, 10 xfailed
  • make check-format: passed
  • uv run mypy .: no issues in 431 source files
  • make docstring-test: 17 passed, 24 skipped
  • make acceptance-test: 199 passed, 72 skipped, 11 deselected
  • make integration-test: 1,016 passed; the two google/gemma-2-2b-it cases could not run locally because this machine is not authenticated for that gated Hugging Face model (401)

Checklist

  • The change is self-explanatory and does not need additional inline comments
  • No documentation change is needed for this bridge-native hook mapping
  • My changes generate no new warnings
  • I have added tests that prove the feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4

Copy link
Copy Markdown
Collaborator

Looks great, thanks @LarryHu0217!

@jlarson4
jlarson4 merged commit 3350142 into TransformerLensOrg:dev-4.x Aug 13, 2026
25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Proposal] Map BERT token-type embeddings so hook_token_type_embed is observable on the bridge

2 participants