fix a crash when using context retrievers

Author: Aethor

renard/pipeline/ner/ner.py

@@ -281,7 +281,7 @@ def batch_labels(
         batch_labels = ["O"] * len(tokens)
 
         try:
-            inference_start = ctxmask[batch_i].tolist().index(1)
+            inference_start = ctxmask[batch_i].tolist().index(0)
         except ValueError:
             inference_start = 0
 
@@ -290,7 +290,7 @@ def batch_labels(
             if token_i is None:
                 continue
 
-            if ctxmask[batch_i][token_i] == 0:
+            if ctxmask[batch_i][token_i] == 1:
                 continue
 
             batch_labels[token_i - inference_start] = wp_label

renard/pipeline/ner/retrieval.py

@@ -66,14 +66,15 @@ def __call__(self, dataset: NERDataset) -> NERDataset:
 
             elements_with_context.append((lctx, elt, rctx))
 
-        return NERDataset(
+        ner_dataset = NERDataset(
             [lctx + element + rctx for lctx, element, rctx in elements_with_context],
             dataset.tokenizer,
             [
                 [1] * len(lctx) + [0] * len(element) + [1] * len(rctx)
                 for lctx, element, rctx in elements_with_context
             ],
         )
+        return ner_dataset
 
 
 class NERSamenounContextRetriever(NERContextRetriever):

tests/test_ner.py

@@ -5,8 +5,9 @@
 from hypothesis.control import assume
 from hypothesis.strategies import lists, sampled_from
 from transformers import BertTokenizerFast
+from renard.pipeline.progress import get_progress_reporter
 from renard.ner_utils import NERDataset
-from renard.pipeline.ner import ner_entities, score_ner
+from renard.pipeline.ner import ner_entities, score_ner, BertNamedEntityRecognizer
 from renard.pipeline.ner.retrieval import (
     NERBM25ContextRetriever,
     NERContextRetriever,
@@ -34,6 +35,22 @@ def test_has_correct_number_of_entities(tokens: List[str]):
     assert len(entities) == len(tokens)
 
 
+@pytest.mark.skipif(os.getenv("RENARD_TEST_SLOW") != "1", reason="performance")
+def test_run_with_context_retriever():
+    ner_step = BertNamedEntityRecognizer(
+        context_retriever=NERNeighborsContextRetriever(k=2)
+    )
+    ner_step._pipeline_init_(lang="eng", progress_reporter=get_progress_reporter(None))
+    # known crash in Renard==0.7.1
+    sentences = [
+        "Whether i shall turn out to be the hero of my own life , or whether that station will be held by anybody else , these pages must show .".split(),
+        "To begin my life with the beginning of my life , i record that i was born ( as i have been informed and believe ) on a friday , at twelve o'clock at night .".split(),
+        "This was the fault of Dr. Strange .".split(),
+    ]
+    tokens = [token for tokens in sentences for token in tokens]
+    _ = ner_step(tokens, sentences)
+
+
 @pytest.mark.skipif(os.getenv("RENARD_TEST_SLOW") != "1", reason="performance")
 @pytest.mark.parametrize(
     "retriever_class", [NERSamenounContextRetriever, NERBM25ContextRetriever]