Add tests for list filter and Gopher quality filter

pyproject.toml

@@ -41,7 +41,7 @@ io = [
   "pyarrow",
   "python-magic",
   "warcio",
-  "datasets"
+  "datasets>=2.18.0"
 ]
 s3 = [
   "s3fs>=2023.12.2",
@@ -85,6 +85,7 @@ merge_stats = "datatrove.tools.merge_stats:main"
 launch_pickled_pipeline = "datatrove.tools.launch_pickled_pipeline:main"
 failed_logs = "datatrove.tools.failed_logs:main"
 inspect_data = "datatrove.tools.inspect_data:main"
+jobs_status = "datatrove.tools.jobs_status:main"
 
 [build-system]
 requires = ["setuptools"]

src/datatrove/executor/slurm.py

@@ -292,8 +292,7 @@ def get_sbatch_args(self, max_array: int = 1) -> dict:
             "array": f"0-{max_array - 1}{f'%{self.workers}' if self.workers != -1 else ''}",
             "requeue": "",
             "qos": self.qos,
-            "mail-type": self.mail_type,
-            "mail-user": self.mail_user,
+            **({"mail-type": self.mail_type, "mail-user": self.mail_user} if self.mail_user else {}),
             **self._sbatch_args,
         }
 

src/datatrove/io.py

@@ -139,11 +139,11 @@ def list_files(
             [
                 f
                 for f, info in (
-                    self.find(subdirectory, maxdepth=0 if not recursive else None, detail=True, **extra_options)
+                    self.find(subdirectory, maxdepth=1 if not recursive else None, detail=True, **extra_options)
                     if not glob_pattern
                     else self.glob(
                         self.fs.sep.join([glob_pattern, subdirectory]),
-                        maxdepth=0 if not recursive else None,
+                        maxdepth=1 if not recursive else None,
                         detail=True,
                         **extra_options,
                     )

src/datatrove/pipeline/readers/huggingface.py

@@ -52,7 +52,7 @@ def get_document_from_dict(self, data: dict, source: str, id_in_file: int | str)
         return document
 
     def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
-        from datasets import load_dataset
+        from datasets import load_dataset  # type: ignore
 
         if data:
             yield from data

src/datatrove/pipeline/readers/warc.py

@@ -28,7 +28,7 @@ class WarcReader(BaseDiskReader):
     """
 
     name = "🕷 Warc"
-    _requires_dependencies = ["warcio", ("cchardet", "faust-chardet"), ("magic", "python-magic")]
+    _requires_dependencies = ["warcio", ("cchardet", "faust-cchardet"), ("magic", "python-magic")]
 
     def __init__(
         self,

src/datatrove/pipeline/tokens/tokenizer.py

@@ -105,6 +105,8 @@ def cleanup(self):
         self.output_folder.rm_file(self.filename)
         if self.loss_file:
             self.output_folder.rm_file(f"{self.filename}.loss")
+        if self.save_final_metadata and self.output_folder.exists(f"{self.filename}.metadata"):
+            self.output_folder.rm_file(f"{self.filename}.metadata")
 
     def write_bytes(self, tk_bytes: bytes, doc_ends: list[int] = None):
         """Write tk_bytes to the tokens file and update the document boundaries with a new document end (in tokens).

src/datatrove/tools/jobs_status.py

@@ -0,0 +1,91 @@
+import argparse
+import json
+import os.path
+
+from loguru import logger
+from rich.console import Console
+
+from datatrove.io import get_datafolder
+from datatrove.utils._import_utils import is_rich_available
+
+
+if not is_rich_available():
+    raise ImportError("Please install `rich` to run this command (`pip install rich`).")
+
+
+parser = argparse.ArgumentParser("Fetch all jobs that are running or complete.")
+
+parser.add_argument(
+    "path", type=str, nargs="?", help="Path to the logging folder. Defaults to current directory.", default=os.getcwd()
+)
+
+parser.add_argument(
+    "-p", "--log_prefix", type=str, nargs="?", help="Prefix of logging folders to be scanned.", default=""
+)
+parser.add_argument("-hc", "--hide_complete", help="Hide all jobs that are already complete.", action="store_true")
+
+
+def main():
+    """
+    Takes a `path` as input, gets all valid job folders and their total number of tasks from `executor.json` and then gets which ranks are
+    incomplete by scanning `path/{LOGGING_DIRS}/completions`. If a `log_prefix` is provided the directories following the `path/log_prefix{LOGGING_DIRS}/completions`
+    pattern are scanned.
+    """
+    args = parser.parse_args()
+    console = Console()
+
+    main_folder = get_datafolder(args.path)
+    logging_dirs = [
+        f
+        for f, info in main_folder.glob(f"{args.log_prefix}*", detail=True, maxdepth=1).items()
+        if info["type"] == "directory"
+    ]
+    logger.remove()
+
+    complete_jobs = 0
+    incomplete_jobs = 0
+
+    for path in logging_dirs:
+        logging_dir = get_datafolder(main_folder.resolve_paths(path))
+        if not logging_dir.isfile("executor.json"):
+            console.log(
+                f'Could not find "executor.json" in the given directory ({path}). Are you sure it is a '
+                "logging folder?",
+                style="red",
+            )
+            continue
+        with logging_dir.open("executor.json", "rt") as f:
+            world_size = json.load(f).get("world_size", None)
+        if not world_size:
+            console.log(
+                f"Could not get the total number of tasks in {path}, please try relaunching the run.",
+                style="red",
+            )
+            continue
+
+        with console.status("Fetching list of incomplete tasks"):
+            completed = set(logging_dir.list_files("completions"))
+            incomplete = set(filter(lambda rank: f"completions/{rank:05d}" not in completed, range(world_size)))
+
+        if len(incomplete) == 0:
+            emoji = "✅"
+            complete_jobs += 1
+        else:
+            emoji = "❌"
+            incomplete_jobs += 1
+
+        if len(incomplete) > 0 or not args.hide_complete:
+            console.log(
+                f"{emoji} {path + ':': <50}{len(completed)}/{world_size} ({len(completed)/(world_size):.0%}) completed tasks."
+            )
+
+    if complete_jobs + incomplete_jobs > 0:
+        console.log(
+            f"Summary: {complete_jobs}/{complete_jobs+incomplete_jobs} ({complete_jobs/(complete_jobs+incomplete_jobs):.0%}) jobs completed."
+        )
+    else:
+        console.log("No jobs found.")
+
+
+if __name__ == "__main__":
+    main()

tests/pipeline/test_filters.py

@@ -6,6 +6,7 @@
     GopherRepetitionFilter,
     LambdaFilter,
     LanguageFilter,
+    ListFilter,
     RegexFilter,
     UnigramLogProbFilter,
     URLFilter,
@@ -65,12 +66,20 @@ def test_gopher_repetition(self):
     def test_gopher_quality(self):
         gopher_quality = GopherQualityFilter(min_doc_words=10, max_doc_words=1000)
         self.check_filter(gopher_quality, get_doc("I am too small..."), "gopher_short_doc")
+        self.check_filter(gopher_quality, get_doc("Very long document. " * 400), "gopher_long_doc")
         self.check_filter(gopher_quality, get_doc("I am " * 20), "gopher_below_avg_threshold")
         self.check_filter(gopher_quality, get_doc("interconnection " * 20), "gopher_above_avg_threshold")
         self.check_filter(gopher_quality, get_doc("# comment " * 20), "gopher_too_many_hashes")
         self.check_filter(gopher_quality, get_doc("... comment " * 20), "gopher_too_many_ellipsis")
+        self.check_filter(gopher_quality, get_doc("• comment\n" * 20), "gopher_too_many_bullets")
+        self.check_filter(
+            gopher_quality,
+            get_doc("comment comment comment comment comment comment comment comment comment...\n" * 20),
+            "gopher_too_many_end_ellipsis",
+        )
         text = "the ./!*?<><> apple <?////> orange  ++ interconnection !<>??? have" * 20
         self.check_filter(gopher_quality, get_doc(text), "gopher_below_alpha_threshold")
+        self.check_filter(gopher_quality, get_doc("No stopwords. " * 10), "gopher_enough_stop_words")
         self.assertTrue(gopher_quality(get_doc(TEXT_LF_1)))
 
     def test_lambda(self):
@@ -119,3 +128,10 @@ def test_url(self):
                 assert url_filter.filter(doc)
             else:
                 self.check_filter(url_filter, doc, result)
+
+    def test_list(self):
+        list_filter = ListFilter()
+        self.check_filter(list_filter, get_doc("List\n" * 5), "Suspected list")
+        self.check_filter(list_filter, get_doc("Also list\n" * 5), "Suspected list")
+        self.check_filter(list_filter, get_doc("And another list\n" * 5), "Suspected list")
+        self.assertTrue(list_filter.filter(get_doc("Not a list anymore\n" * 5)) is True)

tests/pipeline/test_hf_reader.py

@@ -1,14 +1,15 @@
 import unittest
 
+from datatrove.pipeline.readers import HuggingFaceDatasetReader
+
 from ..utils import require_datasets
 
 
 @require_datasets
 class TestHuggingFaceReader(unittest.TestCase):
     def test_read_dataset(self):
-        # reader = HuggingFaceDatasetReader(
-        #     "truthful_qa", dataset_options={"name": "generation", "split": "validation"}, text_key="question"
-        # )
-        # data = list(reader())
-        # assert len(data) == 817
-        pass
+        reader = HuggingFaceDatasetReader(
+            "truthful_qa", dataset_options={"name": "generation", "split": "validation"}, text_key="question"
+        )
+        data = list(reader())
+        assert len(data) == 817