[Q] Supporting multi-file / SentenceTransformers model checkpoints (e.g. EmbeddingGemma) #471

Context

I am working on adding support for EmbeddingGemma (google/embeddinggemma-300m) in PR #469

In terms of architecture, EmbeddingGemma consists of:
1. Gemma 3 backbone
2. Attention-masked mean pooling
3. Two dense linear bottleneck projections ($768 \to 3072 \to 768$, without bias)
4. L2 Normalization

The Problem: Multi-Directory Parameter Checkpoints with Colliding Keys

In standard Hugging Face Transformer models, all weights reside in a single model.safetensors (or a sharded index where tensor names are globally unique).

However, SentenceTransformers models like EmbeddingGemma store modular components in separate files across subdirectories:
* model.safetensors (in root): contains the Gemma 3 backbone.
* 2_Dense/model.safetensors: contains the 1st bottleneck projection (linear.weight).
* 3_Dense/model.safetensors: contains the 2nd bottleneck projection (linear.weight).

This introduces two challenges:
1. Identical layer names / key collisions: Both dense layers use the exact same internal PyTorch parameter name (linear.weight). If loaded and merged into a single state dict naively, one overwrites the other.
2. Repository structure variations: Official repositories (like google/embeddinggemma-300m) and mirrors (like unsloth/embeddinggemma-300m) use the 3-file directory structure defined by modules.json, but consolidated or single-file export formats might also exist in the community. Ideally, Bumblebee could handle both multi-file directory structures and unified root checkpoints.

Current Prototype in PR #469

In PR #469, I implemented a working proof-of-concept:
1. Added an optional extra_params_files/2 callback in Bumblebee.Text.EmbeddingGemma:

       def extra_params_files(_spec, repo_files) do
         for dir <- ["2_Dense", "3_Dense"],
             filename = "#{dir}/model.safetensors",
             Map.has_key?(repo_files, dir) or Map.has_key?(repo_files, filename) do
           {"#{dir}.", filename}
         end
       end

2. Extended Bumblebee.Conversion.PyTorchParams with a :path_prefixes option (%{path => prefix}) to namespace tensors when loading weights from auxiliary files before merging.
3. Allowed params_mapping to support either the prefixed keys ("2_Dense.linear", "3_Dense.linear") or unified keys ("dense_1.linear", "dense_2.linear").

While this works and passes verification against PyTorch output, it is relatively ad-hoc.

Questions & Design Options for Maintainers

Before finalizing the PR, I'd like to get your input on the preferred design:

1. Option A: General SentenceTransformers / modules.json support
   • Inspect modules.json (if present) when downloading parameters, and automatically namespace checkpoints loaded from module subdirectories.
2. Option B: Formal ModelSpec callback for multi-part checkpoints
   • Formalize a callback (e.g., extra_params_files/2 or composite parameter specification) so any model requiring supplementary checkpoint files can declare them and their namespace prefixes cleanly.
3. Option C: Custom parameter loader callback
   • Allow models to define a custom parameters loading pipeline when standard single-file/sharded loading is insufficient.