No constructor-based way to set query_length/document_length on `MultiVectorEncoder`

No constructor-based way to set query_length/document_length on MultiVectorEncoder

Description

Transformer.query_length / Transformer.document_length control the per-task truncation length used by MultiVectorEncoder.encode_query / encode_document. There is no way to set these at MultiVectorEncoder(...) construction time when loading a pretrained checkpoint — every kwarg bucket exposed by __init__ either does nothing or raises. The only thing that actually works is loading the model and then mutating the submodule attribute directly (model[0].query_length = ...), which requires knowing the internal module index/layout of the pipeline and can't be done in one call.

Repro script

from sentence_transformers import MultiVectorEncoder

MODEL = "lightonai/LateOn"
QUERY = "What is the capital city of the country of France in Europe today"
DOC = (
    "The city of Paris has served as the capital of France for many centuries, "
    "and it is well known for the Eiffel Tower, the Louvre museum, and the river Seine "
    "which flows through the heart of the city."
)
Q, D = 5, 9  # strange, obviously-non-default values


def show(label, model):
    t = model[0]
    q = model.encode_query([QUERY])[0]
    d = model.encode_document([DOC])[0]
    print(f"{label}:")
    print(f"  module.query_length={t.query_length}  module.document_length={t.document_length}")
    print(f"  encode_query shape={tuple(q.shape)}  encode_document shape={tuple(d.shape)}")


print("=== Baseline (no override) ===")
model = MultiVectorEncoder(MODEL)
show("baseline", model)
print()

print('=== Attempt: top-level kwargs MultiVectorEncoder(MODEL, query_length=5, document_length=9) ===')
try:
    MultiVectorEncoder(MODEL, query_length=Q, document_length=D)
    print("  constructed (unexpected)")
except TypeError as e:
    print(f"  TypeError: {e}")
print()

print(f'=== Attempt: processor_kwargs={{"query_length": {Q}, "document_length": {D}}} ===')
m = MultiVectorEncoder(MODEL, processor_kwargs={"query_length": Q, "document_length": D})
show("processor_kwargs", m)
print()

print(f'=== Attempt: model_kwargs={{"query_length": {Q}, "document_length": {D}}} ===')
try:
    m = MultiVectorEncoder(MODEL, model_kwargs={"query_length": Q, "document_length": D})
    show("model_kwargs", m)
except Exception as e:
    print(f"  {type(e).__name__}: {e}")
print()

print(f'=== Attempt: config_kwargs={{"query_length": {Q}, "document_length": {D}}} ===')
try:
    m = MultiVectorEncoder(MODEL, config_kwargs={"query_length": Q, "document_length": D})
    show("config_kwargs", m)
except Exception as e:
    print(f"  {type(e).__name__}: {e}")
print()

print("=== Only thing that works: model[0].query_length = 5; model[0].document_length = 9 (post-construction) ===")
model[0].query_length = Q
model[0].document_length = D
show("attribute mutation after load", model)

Actual output

=== Baseline (no override) ===
baseline:
  module.query_length=32  module.document_length=300
  encode_query shape=(16, 128)  encode_document shape=(45, 128)

=== Attempt: top-level kwargs MultiVectorEncoder(MODEL, query_length=5, document_length=9) ===
  TypeError: MultiVectorEncoder.__init__() got an unexpected keyword argument 'query_length'

=== Attempt: processor_kwargs={"query_length": 5, "document_length": 9} ===
processor_kwargs:
  module.query_length=32  module.document_length=300
  encode_query shape=(16, 128)  encode_document shape=(45, 128)

=== Attempt: model_kwargs={"query_length": 5, "document_length": 9} ===
  TypeError: ModernBertModel.__init__() got an unexpected keyword argument 'query_length'

=== Attempt: config_kwargs={"query_length": 5, "document_length": 9} ===
config_kwargs:
  module.query_length=32  module.document_length=300
  encode_query shape=(16, 128)  encode_document shape=(45, 128)

=== Only thing that works: model[0].query_length = 5; model[0].document_length = 9 (post-construction) ===
attribute mutation after load:
  module.query_length=5  module.document_length=9
  encode_query shape=(5, 128)  encode_document shape=(9, 128)

None of the four kwargs exposed by MultiVectorEncoder.__init__ provide a working path:

Expected behavior

Some constructor-time, top-level way to set query_length/document_length when loading a MultiVectorEncoder checkpoint — e.g. MultiVectorEncoder(model_name_or_path, query_length=5, document_length=9) — that overrides the checkpoint's saved lengths without requiring the caller to know the pipeline's internal module layout (model[0]).

Root cause

query_length/document_length are declared Transformer.__init__ params and are listed in Transformer.config_keys (base/modules/transformer.py), so they're persisted in — and loaded straight back from — the Transformer submodule's own saved config.json. They aren't reachable through any of model_kwargs / processor_kwargs / config_kwargs:

• model_kwargs goes to AutoModel.from_pretrained(**model_kwargs), whose target class (e.g. ModernBertModel) has no such param, hence the TypeError.
• config_kwargs goes to AutoConfig.from_pretrained(**config_kwargs), whose PretrainedConfig silently accepts and stores unrecognized kwargs as arbitrary attributes rather than raising.
• processor_kwargs goes to AutoProcessor.from_pretrained(**processor_kwargs), whose PreTrainedTokenizerBase.from_pretrained likewise accepts and silently discards unrecognized kwargs.

Tracing the actual load path for MultiVectorEncoder(model_name_or_path):

1. MultiVectorEncoder.__init__ → BaseModel.__init__ → BaseModel._load_modules (base/model.py), which reads modules.json and, for the Transformer submodule, calls Transformer.load(model_name_or_path, ..., model_kwargs=..., processor_kwargs=..., config_kwargs=..., init_defaults=...).
2. Transformer.load (base/modules/transformer.py) calls Transformer._load_init_kwargs, which loads the submodule's saved config.json (this is where the checkpoint's query_length: 32, document_length: 300 live) and merges in model_kwargs / processor_kwargs / config_kwargs — but only inside those three named buckets, never into a top-level query_length / document_length key.
3. Transformer.load then applies init_defaults with setdefault priority (saved config always wins) and calls cls(model_name_or_path=..., **init_kwargs) — this is where Transformer.__init__ actually runs, setting self.query_length / self.document_length from whatever came out of the saved config file in step 2.

There is currently no caller-override channel for these two scalars anywhere in that chain — only the three dict buckets (none of which reach them) and init_defaults/ _get_module_init_defaults (which is designed for defaults, i.e. it loses to the saved config, not the other way around — it exists for legacy-checkpoint fallback values, e.g. in MultiVectorEncoder._get_module_init_defaults).

Suggested minimal fix

Add a second, symmetric hook alongside the existing init_defaults one, but with override (not default) priority, and expose query_length/document_length as real top-level MultiVectorEncoder.__init__ kwargs that feed it:

1. base/modules/module.py, Module.load() — accept init_overrides: dict | None = None and apply it after the existing init_defaults loop, with plain assignment instead of setdefault:

   python复制

   for key, value in (init_defaults or {}).items():
       config.setdefault(key, value)
   for key, value in (init_overrides or {}).items():
       config[key] = value          # NEW: caller wins over saved config
   return cls(**config)

2. base/modules/transformer.py, Transformer.load() — same two-line addition (it has its own load() override, doesn't call the base one).

3. base/model.py, BaseModel._load_modules() — mirror the existing _get_module_init_defaults wiring with a new _get_module_init_overrides hook (default no-op returning {} on BaseModel), and pass its result through as init_overrides=... alongside the existing init_defaults=....

4. multi_vector_encoder/model.py, MultiVectorEncoder.__init__ — add query_length: int | None = None, document_length: int | None = None params, and implement _get_module_init_overrides to return {"query_length": query_length, "document_length": document_length} (None values filtered out) when class_ref resolves to a Transformer subclass — mirroring the existing _get_module_init_defaults override at the same class, just with override instead of default semantics.

This makes MultiVectorEncoder("lightonai/LateOn", query_length=5, document_length=9) work as one would expect, reuses the existing defaults-hook pattern rather than inventing a new mechanism, and leaves model_kwargs/processor_kwargs/config_kwargs semantics completely unchanged (no regression risk there). Happy to put up a PR along these lines if that shape looks reasonable to the maintainers.

Environment

• sentence-transformers == 6.0.1
• transformers == 5.3.0
• torch == 2.11.0+cu130
• Python 3.13.12
• Model: lightonai/LateOn