Block design

fms_dgt/base/block.py

@@ -0,0 +1,136 @@
+# Standard
+from abc import ABC
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+# Third Party
+from datasets import Dataset
+import pandas as pd
+
+BLOCK_ROW_TYPE = Union[Dict, pd.Series]
+BLOCK_INPUT_TYPE = Union[Iterable[BLOCK_ROW_TYPE], pd.DataFrame, Dataset]
+
+
+class BaseBlock(ABC):
+    """Base Class for all Blocks"""
+
+    def __init__(
+        self,
+        name: str = None,
+        arg_fields: List[str] = None,
+        kwarg_fields: List[str] = None,
+        result_field: str = None,
+    ) -> None:
+
+        self._name = name
+
+        # minor type checking
+        if type(arg_fields) == str:
+            arg_fields = [arg_fields]
+        if type(kwarg_fields) == str:
+            kwarg_fields = [kwarg_fields]
+        if type(result_field) == list:
+            assert (
+                len(result_field) == 1
+            ), "Cannot have multiple 'result' fields for {name}"
+            result_field = result_field[0]
+
+        self._arg_fields = arg_fields
+        self._kwarg_fields = kwarg_fields
+        self._result_field = result_field
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def arg_fields(self):
+        return self._arg_fields
+
+    @property
+    def kwarg_fields(self):
+        return self._kwarg_fields
+
+    @property
+    def result_field(self):
+        return self._result_field
+
+    def generate(
+        self,
+        inputs: BLOCK_INPUT_TYPE,
+        arg_fields: Optional[List[str]] = None,
+        kwarg_fields: Optional[List[str]] = None,
+        result_field: Optional[str] = None,
+    ):
+        raise NotImplementedError
+
+
+class BaseUtilityBlock(BaseBlock):
+    pass
+
+
+class BaseGeneratorBlock(BaseBlock):
+    pass
+
+
+class BaseValidatorBlock(BaseBlock):
+    def __init__(self, filter: bool = False, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._filter_invalids = filter
+
+    def generate(
+        self,
+        inputs: BLOCK_INPUT_TYPE,
+        arg_fields: Optional[List[str]] = None,
+        kwarg_fields: Optional[List[str]] = None,
+        result_field: Optional[List[str]] = None,
+    ):
+        outputs = []
+        for x in inputs:
+            inp_args, inp_kwargs = get_args_kwargs(
+                x, arg_fields or self.arg_fields, kwarg_fields or self.kwarg_fields
+            )
+            res = self._validate(*inp_args, **inp_kwargs)
+            if res or not self._filter_invalids:
+                write_result(x, res, result_field or self.result_field)
+                outputs.append(x)
+        return outputs
+
+    def _validate(self, *args: Any, **kwargs: Any) -> bool:
+        raise NotImplementedError
+
+
+def get_args_kwargs(
+    inp: BLOCK_ROW_TYPE,
+    arg_fields: Optional[List[str]] = None,
+    kwarg_fields: Optional[List[str]] = None,
+):
+    if arg_fields is None:
+        arg_fields = []
+    if kwarg_fields is None:
+        kwarg_fields = []
+
+    if type(inp) == dict:
+        args = [inp.get(arg) for arg in arg_fields]
+        kwargs = {kwarg: inp.get(kwarg) for kwarg in kwarg_fields}
+    elif type(inp) in [pd.DataFrame, Dataset]:
+        args = [inp.get(arg) for arg in arg_fields]
+        kwargs = {kwarg: inp.get(kwarg) for kwarg in kwarg_fields}
+    else:
+        raise ValueError(f"Unexpected input type: {type(inp)}")
+
+    return args, kwargs
+
+
+def write_result(
+    inp: BLOCK_ROW_TYPE,
+    res: Any,
+    result_field: str,
+):
+    assert result_field is not None, "Result field cannot be None!"
+
+    if type(inp) == dict:
+        inp[result_field] = res
+    elif type(inp) in [pd.DataFrame, Dataset]:
+        inp[result_field] = res
+    else:
+        raise ValueError(f"Unexpected input type: {type(inp)}")

fms_dgt/base/databuilder.py

@@ -9,32 +9,26 @@
 # Third Party
 from tqdm import tqdm
 
-# First Party
-from fms_dgt.base.generator import BaseGenerator
-from fms_dgt.base.registry import get_generator, get_validator
+# Local
+from fms_dgt.base.block import BaseBlock
+from fms_dgt.base.registry import get_block
 from fms_dgt.base.task import SdgData, SdgTask
-from fms_dgt.base.validator import BaseValidator
-from fms_dgt.generators.llm import CachingLM, LMGenerator
+from fms_dgt.blocks.generators.llm import CachingLM, LMGenerator
 from fms_dgt.utils import all_annotations, sdg_logger
 
 
 @dataclass
 class DataBuilderConfig(dict):
     # data builder naming/registry
     name: Optional[str] = None
-    generators: Optional[Union[str, list]] = None
-    validators: Optional[Union[str, list]] = None
+    blocks: Optional[dict] = None
     generation_kwargs: Optional[dict] = None
     metadata: Optional[
         dict
     ] = None  # by default, not used in the code. allows for users to pass arbitrary info to data builders
 
     def __post_init__(self) -> None:
-        if self.generation_kwargs is not None:
-            if "temperature" in self.generation_kwargs:
-                self.generation_kwargs["temperature"] = float(
-                    self.generation_kwargs["temperature"]
-                )
+        pass
 
 
 TYPE_KEY = "type"
@@ -66,7 +60,7 @@ def __init__(
         self._restart_generation = restart_generation
 
         # initializing generators / validators
-        self._init_gv(lm_cache=lm_cache)
+        self._init_blocks(lm_cache=lm_cache)
 
         # TODO: Data loader goes here
         self._tasks: List[SdgTask] = [
@@ -91,68 +85,52 @@ def config(self) -> DataBuilderConfig:
         return self._config
 
     @property
-    def generators(self) -> List[BaseGenerator]:
-        """Returns the generators associated with this class."""
-        return self._generators
+    def blocks(self) -> List[BaseBlock]:
+        """Returns the blocks associated with this class."""
+        return self._blocks
+
+    def _init_blocks(self, lm_cache: str = None):
+        self._blocks: List[BaseBlock] = []
+
+        # TODO: need to handle nested blocks
+        for obj_name, obj_config in self.config.blocks.items():
+            obj_kwargs = {**obj_config, "name": obj_name}
+            sdg_logger.debug(
+                "Initializing object %s with config %s", obj_name, obj_config
+            )
 
-    @property
-    def validators(self) -> List[BaseValidator]:
-        """Returns the validators associated with this class."""
-        return self._validators
-
-    def _init_gv(self, lm_cache: str = None):
-        _generators = (
-            [self.config.generators]
-            if type(self.config.generators) == str
-            else self.config.generators
-        )
-        _validators = (
-            [self.config.validators]
-            if type(self.config.validators) == str
-            else self.config.validators
-        )
-        self._generators: List[BaseGenerator] = []
-        self._validators: List[BaseValidator] = []
-
-        # TODO: need to handle nested generators / validators
-        for i, info_src in enumerate([_generators, _validators]):
-            # user may not define a generator / validator
-            if info_src is not None:
-                for obj_name, obj_config in info_src.items():
-                    sdg_logger.debug(
-                        "Initializing object %s with config %s", obj_name, obj_config
-                    )
-                    obj = (get_generator if i == 0 else get_validator)(
-                        obj_config[TYPE_KEY]
-                    )(obj_name, obj_config)
-
-                    if lm_cache is not None and isinstance(obj, LMGenerator):
-                        sdg_logger.info(
-                            "Using cache at %s",
-                            lm_cache + "_rank" + str(obj.rank) + ".db",
-                        )
-                        obj = CachingLM(
-                            obj,
-                            lm_cache
-                            # each rank receives a different cache db.
-                            # necessary to avoid multiple writes to cache at once
-                            + f"_model{os.path.split(obj.model_id_or_path)[-1]}_rank{obj.rank}.db",
-                        )
-
-                    type_annotations = all_annotations(type(self))
-                    assert (
-                        obj_name in type_annotations
-                    ), f"Object {obj_name} is missing from definition of DataBuilder {self.__class__}"
-
-                    obj_type = type_annotations[obj_name]
-
-                    # double check types
-                    assert isinstance(obj, obj_type) or (
-                        isinstance(obj, CachingLM) and isinstance(obj.lm, obj_type)
-                    ), f"Type of retrieved object {obj.__class__} for {obj_name} does not match type {obj_type} specified in DataBuilder {self.__class__}"
-
-                    setattr(self, obj_name, obj)
-                    (self._generators if i == 0 else self._validators).append(obj)
+            assert (
+                TYPE_KEY in obj_kwargs
+            ), f"'type' field missing from {obj_name} in data builder config"
+            obj = get_block(obj_kwargs.pop(TYPE_KEY))(**obj_kwargs)
+
+            if lm_cache is not None and isinstance(obj, LMGenerator):
+                sdg_logger.info(
+                    "Using cache at %s",
+                    lm_cache + "_rank" + str(obj.rank) + ".db",
+                )
+                obj = CachingLM(
+                    obj,
+                    lm_cache
+                    # each rank receives a different cache db.
+                    # necessary to avoid multiple writes to cache at once
+                    + f"_model{os.path.split(obj.model_id_or_path)[-1]}_rank{obj.rank}.db",
+                )
+
+            type_annotations = all_annotations(type(self))
+            assert (
+                obj_name in type_annotations
+            ), f"Object {obj_name} is missing from definition of DataBuilder {self.__class__}"
+
+            obj_type = type_annotations[obj_name]
+
+            # double check types
+            assert isinstance(obj, obj_type) or (
+                isinstance(obj, CachingLM) and isinstance(obj.lm, obj_type)
+            ), f"Type of retrieved object {obj.__class__} for {obj_name} does not match type {obj_type} specified in DataBuilder {self.__class__}"
+
+            setattr(self, obj_name, obj)
+            self._blocks.append(obj)
 
     def execute_tasks(self):
         # main entry point to task execution

fms_dgt/base/instance.py

@@ -1,6 +1,6 @@
 # Standard
-from dataclasses import dataclass, field
-from typing import Any, List, Optional
+from dataclasses import dataclass
+from typing import Any, Optional
 
 
 @dataclass