Merge pull request #220 from KennethEnevoldsen/KennethEnevoldsen-patch-2

fix bug in subset
KennethEnevoldsen · Dec 26, 2023 · 49784ea · 49784ea
2 parents 31dda40 + 9027a49
commit 49784ea
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/augmenty/doc/subset.py b/src/augmenty/doc/subset.py
@@ -50,9 +50,9 @@ def paragraph_subset_augmenter_v1(
 
     # Respect entity spans
     while start != 0 and example.y[start].ent_iob_ not in {"O", "B", ""}:
-        start = -1
-    while end < doc_len - 1 and example.y[end + 1].ent_iob_ not in {"O", "B", ""}:
-        end = +1
+        start -= 1
+    while end < doc_len - 1 and example.y[end - 1].ent_iob_ not in {"O", "B", ""}:
+        end += 1
 
     for k in token_anno:
         token_anno[k] = token_anno[k][start:end]