Skip to content

Fix TopkRetriever metadata collation - #2556

Merged
Myhs-phz merged 2 commits into
open-compass:mainfrom
hongleng:fix/1686-topk-retriever-metadata
Jul 28, 2026
Merged

Fix TopkRetriever metadata collation#2556
Myhs-phz merged 2 commits into
open-compass:mainfrom
hongleng:fix/1686-topk-retriever-metadata

Conversation

@hongleng

Copy link
Copy Markdown
Contributor

Motivation

Fixes #1686.

DataCollatorWithPaddingAndCuda currently adds non-tensor metadata to the
BatchEncoding before calling .to(device).

In Transformers 4.46.2, BatchEncoding.to() retains only tensor values.
Therefore, the metadata wrapper is removed during device transfer, and
TopkRetriever.forward() later raises KeyError: 'metadata'.

Modification

  • Keep metadata as its original ordered list instead of wrapping it in
    ListWrapper.
  • Move the padded tensor batch to the target device before attaching metadata.
  • Remove the no-longer-needed ListWrapper.
  • Add a regression test covering both:
    • the legacy behavior that calls .to() on every batch value;
    • the tensor-only behavior used by newer Transformers versions.

The test also verifies that the collator still returns a single
BatchEncoding, tensor fields remain unchanged, and metadata content and
ordering are preserved.

BC-breaking (Optional)

No backward-incompatible changes are introduced.

The collator continues to return a single BatchEncoding, and
TopkRetriever.forward() continues to require the metadata field. The fix
has been verified with Transformers 4.41.0, 4.46.2, and 4.57.6.

Tests

  • pytest tests/openicl/test_icl_topk_retriever.py -q
    • 1 passed, 2 subtests passed
  • pytest tests/openicl -q
    • 53 passed, 2 subtests passed
  • flake8, isort, yapf, codespell, and the other applicable pre-commit hooks
    passed for the modified files

Checklist

Before PR:

  • Pre-commit or other linting tools are used to fix the potential lint issues.
  • Bug fixes are fully covered by unit tests, and the case that causes the bug is included.
  • The modification is covered by unit tests.
  • Documentation update is not required because this does not change the public API or user-facing behavior.

After PR:

  • Compatibility has been checked with Transformers 4.41.0, 4.46.2, and 4.57.6.
  • CLA has been signed and all committers have signed the CLA in this PR.

@Myhs-phz
Myhs-phz merged commit 8d474f8 into open-compass:main Jul 28, 2026
12 checks passed
@hongleng
hongleng deleted the fix/1686-topk-retriever-metadata branch July 30, 2026 18:34
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.

[Bug] The evaluation of configurations requiring a TopKRetriever, such as for flores_datasets, failed with a KeyError: 'metadata'.

4 participants