Fix medlink collate_fn: robust to heterogeneous per-sample keys - #1222
Merged
Conversation
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.
Contributor
|
looks good. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.