Enforce Dataframe Backend Checks

[ryantwolf]

Description

Refactors all modules to inherit from nemo_curator.modules.base.Module. With this change, the base class now checks to see if the dataframe backend for the Module matches what is passed in, and will throw an error prior to computation if there is a mismatch. We also prove a new module ToBackend() that allows easy chaining of CPU and GPU modules.

Usage

from nemo_curator import Module, Sequential, ToBackend
from nemo_curator.datasets import DocumentDataset

class CPUModule(Module):
    def __init__(self):
        super().__init__(input_backend="pandas")

    def call(self, dataset: DocumentDataset):
        dataset.df["cpu_lengths"] = dataset.df["text"].str.len()
        return dataset


class GPUModule(Module):
    def __init__(self):
        super().__init__(input_backend="cudf")

    def call(self, dataset: DocumentDataset):
        dataset.df["gpu_lengths"] = dataset.df["text"].str.len()
        return dataset

def main():
  client = get_client(cluster_type="gpu")
  # Load a dataset with a pandas CPU backend
  dataset = DocumentDataset.from_pandas(...)
  pipeline = Sequential([
    CPUModule(),
    ToBackend("cudf"),
    GPUModule(),
  ])

  curated_dataset = pipeline(dataset)
  curated_dataset.to_json(...)

if __name__ == "__main__":
  main()

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

[praateekmahajan]

Can we add a test case for this? And also type hints for ddf which would be dd.DataFrame

[ryantwolf]

What part of this do you want to test? I'm not quite seeing which part of this you want me to isolate.

[praateekmahajan]

As someone reading this, it's not obviously clear what will be the output of .__module__.split(".")[0] of the type of object's _meta attribute. So I would've proposed that you move this to a function and have tests for that and use that function here.
However I see you've changed this to is_cudf_type(..) which is much more readable, and now we can unittest that directly that.

[praateekmahajan]

This looks great. Left a review, with mostly nits around type hints and function renaming and some class abstraction opinions.

[ryantwolf]

@sarahyurick @praateekmahajan thanks for your reviews. I addressed the feedback that I immediately agreed with/understood. I left some follow up comments and questions on the other points.

[sarahyurick]

LGTM!

[praateekmahajan]

LGTM thanks for adding this!