use dataset.map in pipeline

Author: ArneBinder

src/pytorch_ie/pipeline.py

@@ -68,6 +68,7 @@ def __init__(
             self._dataloader_params,
             self._forward_params,
             self._postprocess_params,
+            self._dataset_map_params,
         ) = self._sanitize_parameters(**kwargs)
 
     def save_pretrained(self, save_directory: str):
@@ -161,7 +162,7 @@ def _ensure_tensor_on_device(self, inputs, device):
 
     def _sanitize_parameters(
         self, **pipeline_parameters
-    ) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
+    ) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
         """
         _sanitize_parameters will be called with any excessive named arguments from either `__init__` or `__call__`
         methods. It should return 4 dictionaries of the resolved parameters used by the various `preprocess`,
@@ -175,6 +176,7 @@ def _sanitize_parameters(
         dataloader_params = {}
         forward_parameters = {}
         postprocess_parameters: Dict[str, Any] = {}
+        dataset_map_parameters = {}
 
         # set preprocess parameters
         field = pipeline_parameters.get("predict_field")
@@ -196,7 +198,17 @@ def _sanitize_parameters(
             if p_name in pipeline_parameters:
                 postprocess_parameters[p_name] = pipeline_parameters[p_name]
 
-        return preprocess_parameters, dataloader_params, forward_parameters, postprocess_parameters
+        for p_name in ["document_batch_size"]:
+            if p_name in pipeline_parameters:
+                dataset_map_parameters["batch_size"] = pipeline_parameters[p_name]
+
+        return (
+            preprocess_parameters,
+            dataloader_params,
+            forward_parameters,
+            postprocess_parameters,
+            dataset_map_parameters,
+        )
 
     def preprocess(
         self,
@@ -283,27 +295,55 @@ def get_dataloader(
 
         return dataloader
 
+    def _process_documents(
+        self,
+        documents: Sequence[Document],
+        preprocess_params: Dict[str, Any],
+        dataloader_params: Dict[str, Any],
+        forward_params: Dict[str, Any],
+        postprocess_params: Dict[str, Any],
+    ) -> Sequence[Document]:
+        # This creates encodings from the documents. It modifies the documents and may produce multiple entries per
+        # document.
+        model_inputs = self.preprocess(documents, **preprocess_params)
+        # Create a dataloader from the model inputs. This uses taskmodule.collate().
+        dataloader = self.get_dataloader(model_inputs=model_inputs, **dataloader_params)
+
+        show_progress_bar = forward_params.pop("show_progress_bar", False)
+        model_outputs: List = []
+        with torch.no_grad():
+            for batch in tqdm.tqdm(dataloader, desc="inference", disable=not show_progress_bar):
+                output = self.forward(batch, **forward_params)
+                processed_output = self.taskmodule.unbatch_output(output)
+                model_outputs.extend(processed_output)
+
+        assert len(model_inputs) == len(
+            model_outputs
+        ), f"length mismatch: len(model_inputs) [{len(model_inputs)}] != len(model_outputs) [{len(model_outputs)}]"
+
+        documents = self.postprocess(
+            model_inputs=model_inputs,
+            model_outputs=model_outputs,
+            **postprocess_params,
+        )
+        return documents
+
     def __call__(
         self,
         documents: Union[Document, Sequence[Document], Dataset],
         *args,
         **kwargs,
-    ) -> Union[Document, Sequence[Document]]:
+    ) -> Union[Document, Sequence[Document], Dataset]:
         if args:
             logger.warning(f"Ignoring args : {args}")
         (
             preprocess_params,
             dataloader_params,
             forward_params,
             postprocess_params,
+            dataset_map_params,
         ) = self._sanitize_parameters(**kwargs)
 
-        in_place: bool = postprocess_params.get("inplace", True)
-        if in_place and isinstance(documents, Dataset):
-            raise InplaceNotSupportedException(
-                "Datasets can't be modified in place. Please set inplace=False."
-            )
-
         if "TOKENIZERS_PARALLELISM" not in os.environ:
             logger.info(
                 "Disabling tokenizer parallelism, we're using DataLoader multithreading already"
@@ -315,6 +355,7 @@ def __call__(
         dataloader_params = {**self._dataloader_params, **dataloader_params}
         forward_params = {**self._forward_params, **forward_params}
         postprocess_params = {**self._postprocess_params, **postprocess_params}
+        dataset_map_params = {**self._dataset_map_params, **dataset_map_params}
 
         self.call_count += 1
         if self.call_count > 10 and self.device.type == "cuda":
@@ -328,30 +369,39 @@ def __call__(
             single_document = True
             documents = [documents]
 
-        # This creates encodings from the documents. It modifies the documents and may produce multiple entries per
-        # document.
-        model_inputs = self.preprocess(documents, **preprocess_params)
-        # Create a dataloader from the model inputs. This uses taskmodule.collate().
-        dataloader = self.get_dataloader(model_inputs=model_inputs, **dataloader_params)
-
-        show_progress_bar = forward_params.pop("show_progress_bar", False)
-        model_outputs: List = []
-        with torch.no_grad():
-            for batch in tqdm.tqdm(dataloader, desc="inference", disable=not show_progress_bar):
-                output = self.forward(batch, **forward_params)
-                processed_output = self.taskmodule.unbatch_output(output)
-                model_outputs.extend(processed_output)
-
-        assert len(model_inputs) == len(
-            model_outputs
-        ), f"length mismatch: len(model_inputs) [{len(model_inputs)}] != len(model_outputs) [{len(model_outputs)}]"
+        processed_documents: Union[Sequence[Document], Dataset]
+        if isinstance(documents, Dataset):
+            in_place: bool = postprocess_params.get("inplace", True)
+            if in_place:
+                raise InplaceNotSupportedException(
+                    "Datasets can't be modified in place. Please set inplace=False."
+                )
+            # do not show inner progress bar
+            forward_params["show_progress_bar"] = False
+
+            processed_documents = documents.map(
+                self._process_documents,
+                fn_kwargs=dict(
+                    preprocess_params=preprocess_params,
+                    dataloader_params=dataloader_params,
+                    forward_params=forward_params,
+                    postprocess_params=postprocess_params,
+                ),
+                batched=True,
+                **dataset_map_params,
+            )
+        else:
+            processed_documents = self._process_documents(
+                documents=documents,
+                preprocess_params=preprocess_params,
+                dataloader_params=dataloader_params,
+                forward_params=forward_params,
+                postprocess_params=postprocess_params,
+            )
 
-        documents = self.postprocess(
-            model_inputs=model_inputs,
-            model_outputs=model_outputs,
-            **postprocess_params,
-        )
         if single_document:
-            return documents[0]
+            # TODO: fix "type: ignore" (if processed_documents is a Dataset, mypy assumes the result is Dict[Any, Any])
+            processed_document: Document = processed_documents[0]  # type: ignore
+            return processed_document
         else:
-            return documents
+            return processed_documents