Skip to content

fix(retrieval): export deployable Sentence Transformers checkpoints#3021

Open
oliverholworthy wants to merge 8 commits into
mainfrom
oholworthy/fix/preserve-embedding-checkpoint
Open

fix(retrieval): export deployable Sentence Transformers checkpoints#3021
oliverholworthy wants to merge 8 commits into
mainfrom
oholworthy/fix/preserve-embedding-checkpoint

Conversation

@oliverholworthy

@oliverholworthy oliverholworthy commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Make Ministral retrieval fine-tuning produce a stock Hugging Face text checkpoint with generated Sentence Transformers metadata that reflects the trained encoder behavior.

Changelog

  • Load standard ministral3 embedding backbones with the stock Hugging Face class and is_causal=false; retain the custom ministral3_bidirec class only for explicit and legacy checkpoints.
  • Generate Sentence Transformers modules, pooling, normalization, prompts, similarity, and sequence-length metadata from the effective bi-encoder at save time.
  • Restrict standard Sentence Transformers export to representable text encoders and preserve the existing save path for unsupported pooling and multimodal backbones.
  • Reuse the resolved local source revision and tokenizer assets without a second model download, while removing stale multimodal processor metadata from the extracted text checkpoint.
  • Preserve size-based consolidated sharding when an extracted text model no longer matches its parent VLM shard map.
  • Use the native Mistral tokenizer in the Ministral retrieval recipe and accept the collator no-op return_token_type_ids=false argument.

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation? (No separate documentation changes are required; the example recipe is updated.)

If you have not finished some of the above items you can still open "Draft" PR.

Additional Information

  • uv run ruff format --check . passes.
  • uv run ruff check . passes.
  • The focused retrieval suite passes 18 tests, including a CPU regression that exercises the real Transformers FlashAttention adapter and verifies the kernel receives is_causal=false.
  • The consolidated shard-index and size-allocation suite passes 10 tests.
  • A one-step Ministral 3B bi-encoder fine-tune completed on Transformers 5.12.1 using SDPA and produced a consolidated checkpoint.
  • The checkpoint loads locally with stock Ministral3Config / Ministral3Model, is_causal=false, no custom auto_map, and Sentence Transformers 5.4.0 successfully encodes a query to a 3072-dimensional embedding.
  • Source tokenizer and chat-template assets are retained; stale Pixtral processor metadata is intentionally omitted because the exported checkpoint is text-only.
  • No linked issue.

@copy-pr-bot

copy-pr-bot Bot commented Jul 10, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@oliverholworthy oliverholworthy force-pushed the oholworthy/fix/preserve-embedding-checkpoint branch from 7995b5d to 4ef6866 Compare July 13, 2026 14:01
Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
@oliverholworthy oliverholworthy force-pushed the oholworthy/fix/preserve-embedding-checkpoint branch from 4ef6866 to a9f4423 Compare July 13, 2026 16:47
@oliverholworthy

Copy link
Copy Markdown
Contributor Author

/ok to test a9f4423

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
@oliverholworthy oliverholworthy force-pushed the oholworthy/fix/preserve-embedding-checkpoint branch from a9f4423 to e736ed3 Compare July 15, 2026 15:03
@oliverholworthy oliverholworthy marked this pull request as ready for review July 15, 2026 15:08
@oliverholworthy oliverholworthy requested review from a team as code owners July 15, 2026 15:08
@oliverholworthy oliverholworthy added the enhancement New feature or request label Jul 15, 2026
@oliverholworthy oliverholworthy force-pushed the oholworthy/fix/preserve-embedding-checkpoint branch from e736ed3 to a7a95de Compare July 15, 2026 21:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant