diff --git a/pyproject.toml b/pyproject.toml index ffd028c66..77cceb6c1 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/datatrove/executor/slurm.py b/src/datatrove/executor/slurm.py index d06fe5820..5da97fb20 100644 --- a/src/datatrove/executor/slurm.py +++ b/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, } diff --git a/src/datatrove/io.py b/src/datatrove/io.py index 5eb377612..8722ea394 100644 --- a/src/datatrove/io.py +++ b/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, ) diff --git a/src/datatrove/pipeline/readers/huggingface.py b/src/datatrove/pipeline/readers/huggingface.py index 36ea7e4a6..fca57cce6 100644 --- a/src/datatrove/pipeline/readers/huggingface.py +++ b/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 diff --git a/src/datatrove/pipeline/readers/warc.py b/src/datatrove/pipeline/readers/warc.py index e165f8ebc..6249c13fb 100644 --- a/src/datatrove/pipeline/readers/warc.py +++ b/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, diff --git a/src/datatrove/pipeline/tokens/tokenizer.py b/src/datatrove/pipeline/tokens/tokenizer.py index a5cd04d76..f226ff9fd 100644 --- a/src/datatrove/pipeline/tokens/tokenizer.py +++ b/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). diff --git a/src/datatrove/tools/jobs_status.py b/src/datatrove/tools/jobs_status.py new file mode 100644 index 000000000..dc730d62b --- /dev/null +++ b/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() diff --git a/tests/pipeline/test_filters.py b/tests/pipeline/test_filters.py index ebbfae270..ff475e5e4 100644 --- a/tests/pipeline/test_filters.py +++ b/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) diff --git a/tests/pipeline/test_hf_reader.py b/tests/pipeline/test_hf_reader.py index 8df74234f..2c4267ba9 100644 --- a/tests/pipeline/test_hf_reader.py +++ b/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