Update checkpointer

docs/sections/how_to_guides/advanced/checkpointing.md

@@ -57,3 +57,33 @@ The final datasets can be found in the following links:
 - Checkpoint dataset: [distilabel-internal-testing/streaming_test_1](https://huggingface.co/datasets/distilabel-internal-testing/streaming_test_1)
 
 - Final distiset: [distilabel-internal-testing/streaming_test](https://huggingface.co/datasets/distilabel-internal-testing/streaming_test)
+
+### Read back the data
+
+In case we want to take a look at a given filename we can take advantage of the `huggingface_hub` library. We will use the `HfFileSystem` to list all the `jsonl` files in the dataset repository, and download onle of them to show how it works:
+
+```python
+from huggingface_hub import HfFileSystem, hf_hub_download
+
+dataset_name = "distilabel-internal-testing/streaming_test_1"
+fs = HfFileSystem()
+filenames = fs.glob(f"datasets/{dataset_name}/**/*.jsonl")
+
+filename = hf_hub_download(repo_id="distilabel-internal-testing/streaming_test_1", filename="config-0/train-00000.jsonl", repo_type="dataset")
+```
+
+The filename will be downloaded to the default cache, and to read the data we can just proceed as with any other jsonlines file:
+
+```python
+import json
+data = []
+
+with open(filename, "r") as f:
+    data = [json.loads(line) for line in f.readlines()]
+
+# [{'a': 1, 'b': 5},
+#  {'a': 2, 'b': 6},
+#  {'a': 3, 'b': 7},
+# ...
+```
+

src/distilabel/steps/checkpointer.py

@@ -101,7 +101,9 @@ def load(self) -> None:
             self.token = get_hf_token(self.__class__.__name__, "token")
 
         self._api = HfApi(token=self.token)
-        # Create the repo if it doesn't exist
+        self._maybe_create_repo()
+
+    def _maybe_create_repo(self) -> None:
         if not self._api.repo_exists(repo_id=self.repo_id, repo_type="dataset"):
             self._logger.info(f"Creating repo {self.repo_id}")
             self._api.create_repo(
@@ -115,6 +117,8 @@ def process(self, *inputs: StepInput) -> "StepOutput":
                 for item in input:
                     json_line = json.dumps(item, ensure_ascii=False)
                     temp_file.write(json_line + "\n")
+                temp_file.flush()  # Make sure it's written
+                temp_file.seek(0)  # Go back to the beginning
                 try:
                     self._api.upload_file(
                         path_or_fileobj=temp_file.name,

tests/integration/test_checkpointing.py

@@ -0,0 +1,58 @@
+# Copyright 2023-present, Argilla, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+import pytest
+from datasets import Dataset
+
+from distilabel.pipeline import Pipeline
+from distilabel.steps import HuggingFaceHubCheckpointer
+from distilabel.steps.base import Step, StepInput
+
+dataset = Dataset.from_dict({"a": [1, 2] * 50, "b": [5, 6] * 50})
+
+
+if TYPE_CHECKING:
+    from distilabel.typing import StepOutput
+
+
+class DoNothing(Step):
+    def process(self, *inputs: StepInput) -> "StepOutput":
+        for input in inputs:
+            yield input
+
+
+@pytest.mark.skip(reason="Currently cannot obtain the correct HF_TOKEN from the CI")
+def test_checkpointing() -> None:
+    with Pipeline(name="simple-text-generation-pipeline") as pipeline:
+        text_generation = DoNothing(input_batch_size=60)
+        checkpoint = HuggingFaceHubCheckpointer(
+            repo_id="distilabel-internal-testing/__streaming_test_1",
+            private=False,
+            input_batch_size=50,
+        )
+        text_generation >> checkpoint
+    pipeline.run(dataset=dataset, use_cache=False)
+
+    from huggingface_hub import HfFileSystem
+
+    dataset_name = "distilabel-internal-testing/__streaming_test_1"
+    fs = HfFileSystem()
+    filenames = fs.glob(f"datasets/{dataset_name}/**/*.jsonl")
+    assert len(filenames) == 2
+
+
+if __name__ == "__main__":
+    test_checkpointing()