[linjiaye929]

Summary

Fix multi-process aggregation of layerwise reranker scores by concatenating on the pair axis (axis=-1). Flat scalar scores remain one-dimensional and keep the same behavior.

This was found by tracing the multi-device score contract: both LayerWiseLLMReranker and LightweightLLMReranker return layer-major lists for multiple cutoff layers, but AbsReranker.encode_multi_process currently concatenates worker results on axis 0. With two workers:

• Three sentence pairs and two cutoff layers raise ValueError because worker chunks have different lengths.
• Four sentence pairs silently produce shape (4, 2) instead of (2, 4), mixing the layer and pair axes.

The fix preserves layer rows and concatenates each worker's pairs in chunk order. Added 12 parameterized offline tests covering flat versus multi-layer outputs, one/even/uneven input counts, and out-of-order worker completion.

Validation

Validated against upstream 897baa67c0ffddbfccf4e6776ff55310d67809d4, using Python 3.12, CPU-only PyTorch 2.5.1 and Transformers 4.46.0 (the lightweight reranker documents this version).

• New aggregation tests: base 6 failed / 6 passed; fix 12 passed.

• Selected offline aggregation/tokenizer-compatibility tests: 16 passed, 4 deselected.

• Import compatibility tests: 4 passed, 1 skipped (the skipped check is specific to Transformers 5).

• Additional local reproduction uses the actual public reranker classes and actual two-process CPU spawn pools, with tiny randomly initialized upstream MiniCPM/Gemma/BERT models and a locally constructed tokenizer. For both layerwise and lightweight rerankers, three-pair input reproduces the base exception; four-pair input reproduces the incorrect base shape. After the fix, both match their single-process reference, and the ordinary scalar reranker still matches it for both counts.

• Single-cutoff compatibility: both layerwise/lightweight rerankers with three pairs match their single-process reference on both base and fixed versions (four additional real-spawn runs).

No pretrained weights/tokenizers, external inference services, GPU runs, or training were used. The random-model reproduction verifies score aggregation, not pretrained model quality. This is focused validation, not the full model-dependent suite. The existing process-pool shutdown emits a resource-tracker semaphore warning in the local harness on both base and fixed versions; no change to pool lifecycle is included.

AI-assisted implementation and testing; submitted for maintainer review.