Skip to content

Commit

Permalink
Fix make_disjoint_window for tail case
Browse files Browse the repository at this point in the history
When the continuation list contains a single entry,
make_disjoint_window incorrectly truncates the context list to be
empty.

The fix checks for the non-overlapping case and, in that case, simply
returns the existing lists.
  • Loading branch information
richhankins committed Jul 22, 2022
1 parent 13ed134 commit 834de32
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 3 deletions.
4 changes: 1 addition & 3 deletions lm_eval/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,10 +114,8 @@ def get_rolling_token_windows(token_list, prefix_token, max_seq_len, context_len

def make_disjoint_window(pair):
"""Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""

a, b = pair

return a[: -(len(b) - 1)], b
return a[: len(a) - (len(b) - 1)], b


class Reorderer:
Expand Down
1 change: 1 addition & 0 deletions tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,3 +219,4 @@ def test_make_disjoint_window():
[2, 3, 4, 5, 6],
)
assert make_disjoint_window(([1, 2, 3, 4, 5], [4, 5, 6])) == ([1, 2, 3], [4, 5, 6])
assert make_disjoint_window(([1, 2, 3, 4, 5], [6])) == ([1, 2, 3, 4, 5], [6])

0 comments on commit 834de32

Please sign in to comment.