Merge pull request #336 from richhankins/fix_disjoint

Fix make_disjoint_window for tail case

lm_eval/utils.py

@@ -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:

tests/test_utils.py

@@ -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])