Skip to content

Fix medlink collate_fn: robust to heterogeneous per-sample keys - #1222

Merged
jhnwu3 merged 1 commit into
sunlabuiuc:masterfrom
lehendo:medlink-collate-fix
Sep 2, 2026
Merged

Fix medlink collate_fn: robust to heterogeneous per-sample keys#1222
jhnwu3 merged 1 commit into
sunlabuiuc:masterfrom
lehendo:medlink-collate-fix

Conversation

@lehendo

@lehendo lehendo commented Aug 30, 2026

Copy link
Copy Markdown
Collaborator

collate_fn built its output keys from only the first sample in a batch. If a batch mixed samples that had a mined hard negative (s_n) with samples that didn't, this either crashed with KeyError('s_n') or silently produced a shorter, misaligned list for that field.

Fix: union keys across all samples, and drop s_n for the whole batch if any sample in it is missing one.

get_train_dataloader's s_n (hard negative) key is only present on samples
where a hard negative was successfully mined. collate_fn initialized its
output dict's keys from only samples[0], so depending on sample order this
either raised KeyError('s_n') (the first sample lacked it but a later one
had it) or silently produced a shorter, misaligned list for that key (the
first sample had it, a later one didn't) -- both reproduced directly.

MedLink.forward consumes s_n as a whole-batch field (corpus = s_p + s_n),
so a partially-present s_n would corrupt that concatenation rather than
just misalign cleanly. Fixed by computing the union of keys across all
samples (order-preserving via dict.fromkeys, not a set, so doctest output
is deterministic) and dropping s_n for the whole batch if any sample in
it lacks one -- equivalent to training that batch without hard negatives,
a mode the model already explicitly supports via s_n=None.

This bug pre-dates PR sunlabuiuc#1195 (byte-identical on master before any of those
changes) and is orthogonal to what that PR fixes; branched separately per
the maintainer's direction so sunlabuiuc#1195 isn't blocked on this.
@fbonc

fbonc commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

looks good.

@jhnwu3 jhnwu3 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@jhnwu3
jhnwu3 merged commit 3b2b879 into sunlabuiuc:master Sep 2, 2026
2 checks passed
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.

3 participants