):
        """Initialize the TensorflowDatasetLoader.
        Args:
            dataset_name: the name of the dataset to load
            split_name: the name of the split to load.
            load_max_docs: a limit to the number of loaded documents. Defaults to 100.
            sample_to_document_function: a function that converts a dataset sample
                into a Document.
        """
        self.dataset_name: str = dataset_name
        self.split_name: str = split_name
        self.load_max_docs = load_max_docs
        """The maximum number of documents to load."""
        self.sample_to_document_function: Optional[
            Callable[[Dict], Document]
        ] = sample_to_document_function
        """Custom function that transform a dataset sample into a Document."""
        self._tfds_client = TensorflowDatasets(
            dataset_name=self.dataset_name,
            split_name=self.split_name,
            load_max_docs=self.load_max_docs,
            sample_to_document_function=self.sample_to_document_function,
        )
[docs]    def lazy_load(self) -> Iterator[Document]:
        yield from self._tfds_client.lazy_load()
[docs]    def load(self) -> List[Document]:
        return list(self.lazy_load())