Taggers for URL filtering

Cargo.toml

@@ -43,6 +43,7 @@ tokio = {version = "1.27.0", features = ["full"]}
 tokio-util = "0.7.7"
 unicode-segmentation = "1.7"
 openssl = { version = "0.10.63", features = ["vendored"] }
+adblock = { version = "0.8.6", features = ["content-blocking"] }
 
 # [target.'cfg(target_arch = "aarch64")'.dependencies]
 # openssl = { version = "0.10.63", features = ["vendored"] }

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "LTpycld2==0.42",     # fork of pycld2 that works on Apple Silicon
     # "pycld2==0.41",
     # "pycld3==0.22",       # does not install correctly
+    "platformdirs>=4.2.0",
     "pyyaml",
     "requests",
     "rich",

python/dolma/__init__.py

@@ -24,13 +24,32 @@
 
 
 def deduper(config: dict):
+    """
+    Run the deduper with the given configuration.
+
+    Args:
+        config (dict): The configuration for the deduper.
+
+    Raises:
+        DolmaRustPipelineError: If there is an error running the deduper.
+
+    """
     try:
         _dolma.deduper_entrypoint(json.dumps(config))
     except RuntimeError as e:
         raise DolmaRustPipelineError(f"Error running deduper: {e}") from e
 
 
 def mixer(config: dict):
+    """
+    Run the mixer with the given configuration.
+
+    Args:
+        config (dict): A dictionary containing the configuration parameters for the mixer.
+
+    Raises:
+        DolmaRustPipelineError: If an error occurs while running the mixer.
+    """
     try:
         _dolma.mixer_entrypoint(json.dumps(config))
     except RuntimeError as e:

python/dolma/core/loggers.py

@@ -1,5 +1,8 @@
 import logging
 import multiprocessing
+from typing import Union
+
+DOLMA_PREFIX = "dolma"
 
 
 def get_logger(name: str) -> logging.Logger:
@@ -13,7 +16,34 @@ def get_logger(name: str) -> logging.Logger:
 
     if not logger.handlers:
         handler = logging.StreamHandler()
-        handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s"))
+        formatter = logging.Formatter(
+            "[%(asctime)s %(name)s %(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+        )
+        handler.setFormatter(formatter)
         logger.addHandler(handler)
 
     return logger
+
+
+def reset_level(level: Union[int, str]) -> None:
+    """
+    Reset the log level for all Dolma loggers.
+
+    Args:
+        level (Union[int, str]): The log level to set. It can be either an integer
+            representing the log level (e.g., logging.DEBUG) or a string
+            representing the log level name (e.g., 'debug').
+
+    Returns:
+        None
+    """
+    if isinstance(level, str):
+        if (level_tmp := getattr(logging, level.strip().upper(), None)) is not None:
+            level = level_tmp
+        else:
+            raise ValueError(f"Invalid log level: {level}")
+
+    for logger in logging.Logger.manager.loggerDict.values():
+        if isinstance(logger, logging.Logger):
+            if logger.name.startswith("dolma"):
+                logger.setLevel(level)

python/dolma/core/parallel.py

@@ -2,7 +2,6 @@
 import itertools
 import logging
 import multiprocessing
-import os
 import pickle
 import random
 import re
@@ -26,6 +25,7 @@
     join_path,
     make_relative,
     mkdir_p,
+    parent,
     split_path,
     sub_prefix,
 )
@@ -175,6 +175,10 @@ def _process_single_and_save_status(
     ):
         """A wrapper around process single that saves a metadata file if processing is successful."""
 
+        # make destination directory if it doesn't exist for the destination and metadata paths
+        mkdir_p(parent(destination_path))
+        mkdir_p(parent(metadata_path))
+
         kwargs = pickle.loads(serialized_kwargs)
         retries_on_error = kwargs.get("retries_on_error", 0) + 1
         while True:
@@ -188,6 +192,7 @@ def _process_single_and_save_status(
                 if retries_on_error == 0:
                     raise DolmaError from exception
 
+        # write the metadata file
         with smart_open.open(metadata_path, "wt") as f:
             f.write(datetime.now().isoformat())
 
@@ -368,13 +373,6 @@ def _get_all_paths(self) -> Tuple[List[str], List[str], List[str]]:
                 if not self._valid_path(path):
                     continue
 
-                # get relative path from source prefix
-                rel_dir, _ = os.path.split(path)
-
-                # make sure destination/metadata directories exists
-                mkdir_p(os.path.join(dst_prefix, rel_dir))
-                mkdir_p(os.path.join(meta_prefix, rel_dir))
-
                 # create new paths to pass to taggers
                 all_source_paths.append(add_suffix(prefix, path))
                 all_destination_paths.append(add_suffix(dst_prefix, path))