Skip to content

Fix TopkRetriever metadata collation#2556

Open
hongleng wants to merge 1 commit into
open-compass:mainfrom
hongleng:fix/1686-topk-retriever-metadata
Open

Fix TopkRetriever metadata collation#2556
hongleng wants to merge 1 commit into
open-compass:mainfrom
hongleng:fix/1686-topk-retriever-metadata

Conversation

@hongleng

Copy link
Copy Markdown

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.

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'.

2 participants