towards being pythonic in slicing for character spans, re: Slicing in…

… SurprisalArray is not fully Pythonic #6
aalok-sathe · Nov 15, 2023 · c8146dd · c8146dd
1 parent 83ca06d
commit c8146dd
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/surprisal/utils.py b/surprisal/utils.py
@@ -21,7 +21,12 @@ def hf_pick_matching_token_ixs(
     """
     span_of_interest = slice(
         span_of_interest.start or 0,
-        span_of_interest.stop or len(encoding.ids),
+        span_of_interest.stop
+        or (
+            len(encoding.ids)
+            if span_type == "word"
+            else encoding.token_to_chars((len(encoding.ids) - 1)).end
+        ),
         span_of_interest.step,
     )