Add the NER retrieval module to reference documentation

Author: Aethor

docs/reference.rst

@@ -64,6 +64,13 @@ NLTKNamedEntityRecognizer
    :members:
 
 
+Retrieval
+---------
+
+.. automodule:: renard.pipeline.ner.retrieval
+   :members:
+
+
 Coreference Resolution
 ======================
 

renard/pipeline/ner/retrieval.py

@@ -29,6 +29,8 @@ def __hash__(self) -> int:
 
 
 class NERContextRetriever:
+    """Base class for NER context retrievers."""
+
     def __init__(self, k: int) -> None:
         self.k = k
 
@@ -278,7 +280,9 @@ def predict(self, examples: List[NERNeuralContextRetrievalExample]) -> torch.Ten
         self.ctx_classifier = self.ctx_classifier.to(self.device)
 
         data_collator = DataCollatorWithPadding(dataset.tokenizer)  # type: ignore
-        dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False, collate_fn=data_collator)  # type: ignore
+        dataloader = DataLoader(
+            dataset, batch_size=self.batch_size, shuffle=False, collate_fn=data_collator
+        )  # type: ignore
 
         # inference using self.ctx_classifier
         self.ctx_classifier = self.ctx_classifier.eval()