fix frame when it is out of bounds

ArneBinder · Oct 29, 2024 · 95ba402 · 95ba402
1 parent 99f022b
commit 95ba402
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/pie_modules/taskmodules/re_text_classification_with_indices.py b/src/pie_modules/taskmodules/re_text_classification_with_indices.py
@@ -769,7 +769,14 @@ def encode_input(
                         for arg in args:
                             arg_center = (arg.token_span.end + arg.token_span.start) // 2
                             arg_frame_start = arg_center - max_tokens_per_argument // 2
+                            # shift the frame to the right if it is out of bounds
+                            if arg_frame_start < 0:
+                                arg_frame_start = 0
                             arg_frame_end = arg_frame_start + max_tokens_per_argument
+                            # shift the frame to the left if it is out of bounds
+                            if arg_frame_end > len(input_ids):
+                                arg_frame_end = len(input_ids)
+                                arg_frame_start = arg_frame_end - max_tokens_per_argument
                             mask[arg_frame_start:arg_frame_end] = 1
                         offsets = np.cumsum(mask != 1)
                         arg_cluster_offset_values = set()