correctly truncates NER context_mask

CompNet · Dec 26, 2024 · 2898194 · 2898194
1 parent 8fdda13
commit 2898194
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/renard/ner_utils.py b/renard/ner_utils.py
@@ -104,9 +104,18 @@ def __getitem__(self, index: int) -> BatchEncoding:
             truncation=True,
             max_length=512,  # TODO
             is_split_into_words=True,
+            return_length=True,
         )
 
-        batch["context_mask"] = self._context_mask[index]
+        length = batch["length"][0]
+        del batch["length"]
+        if self.tokenizer.truncation_side == "right":
+            batch["context_mask"] = self._context_mask[index][:length]
+        else:
+            assert self.tokenizer.truncation_side == "left"
+            batch["context_mask"] = self._context_mask[index][
+                len(batch["input_ids"]) - length :
+            ]
 
         return batch