Many Dolma toolkit functions use dolma.core.parallel.BaseParallelProcessor to parallelize tasks over a list of inputs while tracking progress via progress bars. To create a custom processor, you must subclass BaseParallelProcessor and implement two specific class methods:
process_single(cls, source_path, destination_path, queue, **kwargs): This method contains the core logic. It is called for each individual input file. You are responsible for opening the source_path, processing the data, and writing the results to destination_path. During processing, you should periodically call increment_progressbar to update the status.increment_progressbar(cls, queue, /, ...): This method updates the progress bars. Any arguments provided after the / separator in the signature define the metrics tracked by the progress bars. You must call super().increment_progressbar(...) within this method, passing the same arguments to ensure the base class updates the shared queue correctly.
from dolma.core.parallel import BaseParallelProcessor
from queue import Queue
class CustomParallelProcessor(BaseParallelProcessor):
@classmethod
def increment_progressbar(
cls,
queue: Queue,
/,
files: int = 0,
documents: int = 0,
...
):
"""
This method is called in the process_single
to increment the progress bar. You can create as many progress bars as are
the numbers of arguments after the '/' separator.
"""
super().increment_progressbar(
queue,
files=files,
documents=documents,
...
)
@classmethod
def process_single(
cls,
source_path: str,
destination_path: str,
queue: Queue,
**kwargs: Any,
):
"""
This method is to process a single input file.
The method broadly opens source_path file, processes it and writes the output to
destination_path.
"""
...
