sdg-hub 0.1.0a1__py3-none-any.whl

sdg_hub/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# Local
+from .sdg import SDG
+from ._version import __version__

sdg_hub/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.0a1'
+__version_tuple__ = version_tuple = (0, 1, 0)

sdg_hub/blocks/__init__.py

@@ -0,0 +1,6 @@
+# Local
+from .block import *
+from .filterblock import *
+from .iterblock import *
+from .llmblock import *
+from .utilblocks import *

sdg_hub/blocks/block.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from abc import ABC
+from collections import ChainMap
+from typing import Any, Dict, Union
+
+# Third Party
+from jinja2 import Template, UndefinedError
+import yaml
+
+# Local
+from ..registry import BlockRegistry
+from ..logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+@BlockRegistry.register("Block")
+class Block(ABC):
+    def __init__(self, block_name: str) -> None:
+        self.block_name = block_name
+
+    @staticmethod
+    def _validate(prompt_template: Template, input_dict: Dict[str, Any]) -> bool:
+        """
+        Validate the input data for this block. This method validates whether all required
+        variables in the Jinja template are provided in the input_dict.
+
+        :param prompt_template: The Jinja2 template object.
+        :param input_dict: A dictionary of input values to check against the template.
+        :return: True if the input data is valid (i.e., no missing variables), False otherwise.
+        """
+        
+        class Default(dict):
+            def __missing__(self, key: str) -> None:
+                raise KeyError(key)
+        
+        try:
+            # Try rendering the template with the input_dict
+            prompt_template.render(ChainMap(input_dict, Default()))
+            return True
+        except UndefinedError as e:
+            logger.error(f"Missing key: {e}")
+            return False
+
+    def _load_config(self, config_path: str) -> Union[Dict[str, Any], None]:
+        """
+        Load the configuration file for this block.
+
+        :param config_path: The path to the configuration file.
+        :return: The loaded configuration.
+        """
+        with open(config_path, "r", encoding="utf-8") as config_file:
+            return yaml.safe_load(config_file)

sdg_hub/blocks/filterblock.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import operator
+
+# Third Party
+from datasets import Dataset
+
+# Local
+from .block import Block
+from ..registry import BlockRegistry
+from ..logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+@BlockRegistry.register("FilterByValueBlock")
+class FilterByValueBlock(Block):
+    def __init__(
+        self, filter_column, filter_value, operation, convert_dtype=None, **batch_kwargs
+    ) -> None:
+        """
+        Initializes a new instance of the FilterByValueBlock class.
+
+        Parameters:
+        - filter_column (str): The name of the column in the dataset to apply the filter on.
+        - filter_value (any or list of any): The value(s) to filter by.
+        - operation (callable): A function that takes two arguments (column value and filter value) and returns a boolean indicating whether the row should be included in the filtered dataset.
+        - convert_dtype (callable, optional): A function to convert the data type of the filter column before applying the filter. Defaults to None.
+        - **batch_kwargs: Additional kwargs for batch processing.
+
+        Returns:
+        None
+        """
+        super().__init__(block_name=self.__class__.__name__)
+        self.value = filter_value if isinstance(filter_value, list) else [filter_value]
+        self.column_name = filter_column
+        self.operation = operation
+        self.convert_dtype = convert_dtype
+        self.num_procs = batch_kwargs.get("num_procs", 1)
+
+    def _convert_dtype(self, sample):
+        try:
+            sample[self.column_name] = self.convert_dtype(sample[self.column_name])
+        except ValueError as e:
+            logger.error(
+                "Error converting dtype: %s, filling with None to be filtered later", e
+            )
+            sample[self.column_name] = None
+        return sample
+
+    def generate(self, samples) -> Dataset:
+        if self.convert_dtype:
+            samples = samples.map(
+                self._convert_dtype,
+                num_proc=self.num_procs,
+            )
+
+        if self.operation == operator.contains:
+            samples = samples.filter(
+                lambda x: self.operation(self.value, x[self.column_name]),
+                num_proc=self.num_procs,
+            )
+
+        samples = samples.filter(
+            lambda x: x[self.column_name] is not None,
+            num_proc=self.num_procs,
+        )
+
+        samples = samples.filter(
+            lambda x: any(
+                self.operation(x[self.column_name], value) for value in self.value
+            ),
+            num_proc=self.num_procs,
+        )
+
+        return samples

sdg_hub/blocks/iterblock.py

@@ -0,0 +1,31 @@
+# Third Party
+from datasets import Dataset
+
+# Local
+from .block import Block
+from ..registry import BlockRegistry
+from ..logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+@BlockRegistry.register("IterBlock")
+class IterBlock(Block):
+    def __init__(self, block_name, num_iters, block_type, block_kwargs, **kwargs):
+        super().__init__(block_name)
+        self.num_iters = num_iters
+        self.block = block_type(**block_kwargs)
+        self.gen_kwargs = kwargs.get("gen_kwargs", {})
+        self.gen_kwargs = kwargs.get("gen_kwargs", {})
+
+    def generate(self, samples, **gen_kwargs) -> Dataset:
+        generated_samples = []
+        num_iters = self.num_iters
+
+        for _ in range(num_iters):
+            batch_generated = self.block.generate(
+                samples, **{**self.gen_kwargs, **gen_kwargs}
+            )
+            generated_samples.extend(batch_generated)
+
+        return Dataset.from_list(generated_samples)

sdg_hub/blocks/llmblock.py

@@ -0,0 +1,430 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from collections import Counter
+from typing import Any, Dict, List
+import json
+import re
+
+# Third Party
+from datasets import Dataset
+from jinja2 import Template
+import openai
+
+# Local
+from .block import Block
+from ..logger_config import setup_logger
+from ..registry import BlockRegistry, PromptRegistry
+
+logger = setup_logger(__name__)
+
+
+def server_supports_batched(client, model_id: str) -> bool:
+    supported = getattr(client, "server_supports_batched", None)
+    if supported is not None:
+        return supported
+    try:
+        # Make a test call to the server to determine whether it supports
+        # multiple input prompts per request and also the n parameter
+        response = client.completions.create(
+            model=model_id, prompt=["test1", "test2"], max_tokens=1, n=3
+        )
+        # Number outputs should be 2 * 3 = 6
+        supported = len(response.choices) == 6
+    except openai.InternalServerError:
+        supported = False
+    client.server_supports_batched = supported
+    logger.info(f"LLM server supports batched inputs: {client.server_supports_batched}")
+    return supported
+
+
+@BlockRegistry.register("LLMBlock")
+# pylint: disable=dangerous-default-value
+class LLMBlock(Block):
+    # pylint: disable=too-many-instance-attributes
+    def __init__(
+        self,
+        block_name,
+        config_path,
+        client,
+        output_cols,
+        parser_kwargs={},
+        model_prompt="{prompt}",
+        model_id=None,
+        **batch_kwargs,
+    ) -> None:
+        super().__init__(block_name)
+        self.block_config = self._load_config(config_path)
+        self.prompt_struct = (
+            """{system}\n{introduction}\n{principles}\n{examples}\n{generation}"""
+        )
+        filtered_config = {
+            k: (v if v is not None else "") for k, v in self.block_config.items()
+        }
+        self.prompt_template = Template(self.prompt_struct.format(**filtered_config))
+        self.client = client
+        if model_id:
+            self.model = model_id
+        else:
+            # get the default model id from client
+            self.model = self.client.models.list().data[0].id
+
+        self.model_prompt = model_prompt
+        self.output_cols = output_cols
+        self.batch_params = batch_kwargs.get("batch_kwargs", {})
+        self.parser_name = parser_kwargs.get("parser_name", None)
+        self.parsing_pattern = parser_kwargs.get("parsing_pattern", None)
+        self.parser_cleanup_tags = parser_kwargs.get("parser_cleanup_tags", None)
+        self.defaults = {
+            "model": self.model,
+            "temperature": 0,
+            "max_tokens": 4096,
+        }
+
+        # Whether the LLM server supports a list of input prompts
+        # and supports the n parameter to generate n outputs per input
+        self.server_supports_batched = server_supports_batched(client, self.model)
+
+    def _parse(self, generated_string) -> dict:
+        matches = {}
+
+        if self.parser_name is not None and self.parser_name == "custom":
+            pattern = re.compile(self.parsing_pattern, re.DOTALL)
+            all_matches = pattern.findall(generated_string)
+            matches = {column_name: [] for column_name in self.output_cols}
+            if all_matches and isinstance(all_matches[0], tuple):
+                for match in all_matches:
+                    for column_name, value in zip(self.output_cols, match):
+                        value = value.strip()
+                        for clean_tag in self.parser_cleanup_tags:
+                            value = value.replace(clean_tag, "")
+                        matches[column_name].append(value)
+            else:
+                matches[self.output_cols[0]] = (
+                    [match.strip() for match in all_matches] if all_matches else []
+                )
+        else:
+            for start_tag, end_tag, output_col in zip(
+                self.block_config.get("start_tags", []),
+                self.block_config.get("end_tags", []),
+                self.output_cols,
+            ):
+                if not start_tag and not end_tag:
+                    matches[output_col] = [
+                        generated_string.strip() if generated_string else None
+                    ]
+                else:
+                    pattern = re.escape(start_tag) + r"(.*?)" + re.escape(end_tag)
+                    all_matches = re.findall(pattern, generated_string, re.DOTALL)
+                    matches[output_col] = (
+                        [match.strip() for match in all_matches] if all_matches else []
+                    )
+
+        return matches
+
+    def _format_prompt(self, sample: Dict) -> str:
+        prompt_templated_str = self.prompt_template.render(sample).strip()
+        return PromptRegistry.render_template(
+            self.model_prompt, prompt_templated_str, add_generation_prompt=True
+        ).strip()
+
+    def _generate(self, samples, **gen_kwargs) -> list:
+        prompts = [self._format_prompt(sample) for sample in samples]
+        logger.debug("Prompt: %s", prompts[0])
+        generate_args = {**self.defaults, **gen_kwargs}
+
+        if self.server_supports_batched:
+            response = self.client.completions.create(prompt=prompts, **generate_args)
+            # if stop is provided, then we need to add the stop token to the generated text,
+            # this is because the stop token is not included in the generated text - this is a limitation of the openai api
+            # we need to add the stop token to the generated text to make it consistent for the parser
+            if "stop" in generate_args:
+                return [
+                    choice.text.strip() + "".join(generate_args["stop"])
+                    for choice in response.choices
+                ]
+            return [choice.text.strip() for choice in response.choices]
+
+        n = gen_kwargs.get("n", 1)
+        results = []
+        for prompt in prompts:
+            for _ in range(n):
+                response = self.client.completions.create(
+                    prompt=prompt, **generate_args
+                )
+                if "stop" in generate_args:
+                    results.append(
+                        response.choices[0].text.strip()
+                        + "".join(generate_args["stop"])
+                    )
+                results.append(response.choices[0].text.strip())
+        return results
+
+    def generate(self, samples: Dataset, **gen_kwargs) -> Dataset:
+        """
+        Generate the output from the block. This method should first validate the input data,
+        then generate the output, and finally parse the generated output before returning it.
+
+        :return: The parsed output after generation.
+        """
+        num_samples = self.block_config.get("num_samples", None)
+        logger.debug("Generating outputs for {} samples".format(len(samples)))
+
+        if (num_samples is not None) and ("num_samples" not in samples.column_names):
+            samples = samples.add_column("num_samples", [num_samples] * len(samples))
+
+        # validate each sample
+        # Log errors and remove invalid samples
+        valid_samples = []
+
+        for sample in samples:
+            if self._validate(self.prompt_template, sample):
+                valid_samples.append(sample)
+            else:
+                logger.warning(
+                    f"Sample failed validation: {sample}"
+                )  # Log details of the failed sample
+
+        samples = valid_samples
+
+        if len(samples) == 0:
+            logger.warning(
+                "No valid samples to generate outputs for, returning empty dataset"
+            )
+            return Dataset.from_list([])
+
+        # generate the output
+
+        outputs = self._generate(samples, **gen_kwargs)
+
+        logger.debug("Generated outputs: %s", outputs)
+
+        num_parallel_samples = gen_kwargs.get("n", 1)
+        extended_samples = []
+
+        # Duplicate each input sample n times, where n is the number
+        # of output sequences generated per input, so that we can
+        # pair up the inputs and outputs.
+        for item in samples:
+            extended_samples.extend([item] * num_parallel_samples)
+
+        new_data = []
+        for sample, output in zip(extended_samples, outputs):
+            parsed_outputs = self._parse(output)
+            max_length = max(len(value) for value in parsed_outputs.values())
+            for values in zip(*(lst[:max_length] for lst in parsed_outputs.values())):
+                new_data.append({**sample, **dict(zip(parsed_outputs.keys(), values))})
+
+        return Dataset.from_list(new_data)
+
+
+@BlockRegistry.register("ConditionalLLMBlock")
+class ConditionalLLMBlock(LLMBlock):
+    def __init__(
+        self,
+        block_name,
+        config_paths,
+        client,
+        model_id,
+        output_cols,
+        selector_column_name,
+        model_prompt="{prompt}",
+        **batch_kwargs,
+    ) -> None:
+        super().__init__(
+            block_name=block_name,
+            config_path=list(config_paths.values())[0],
+            client=client,
+            model_id=model_id,
+            output_cols=output_cols,
+            model_prompt=model_prompt,
+            **batch_kwargs,
+        )
+        self.selector_column_name = selector_column_name
+        self.prompt_template = {}
+        if "All" in config_paths:
+            self.prompt_template = self.prompt_struct.format(**self.block_config)
+        else:
+            for config_key, config in config_paths.items():
+                # Template(self.prompt_struct.format(**filtered_config))
+                filtered_config = {
+                    k: (v if v is not None else "") for k, v in self.block_config.items()
+                }
+                self.prompt_template[config_key] = Template(self.prompt_struct.format(
+                    **self._load_config(config)
+                ))
+
+    def _format_prompt(self, sample: Dict) -> str:
+        if isinstance(self.prompt_template, dict):
+            return (
+                self.prompt_template[sample[self.selector_column_name]]
+                .render(**sample)
+                .strip()
+            )
+
+        return self.prompt_template.render(**sample).strip()
+
+    def _validate(self, prompt_template: str, input_dict: Dict[str, Any]) -> bool:
+        if isinstance(prompt_template, dict):
+            prompt_template = prompt_template[input_dict[self.selector_column_name]]
+        return super()._validate(prompt_template, input_dict)
+
+
+@BlockRegistry.register("LLMLogProbBlock")
+class LLMLogProbBlock(LLMBlock):
+    # init with init of the parent class
+    def __init__(
+        self,
+        block_name,
+        config_path,
+        client,
+        output_cols,
+        parser_kwargs={},
+        model_prompt="{prompt}",
+        model_id=None,
+        **batch_kwargs,
+    ) -> None:
+        super().__init__(
+            block_name=block_name,
+            config_path=config_path,
+            client=client,
+            output_cols=output_cols,
+            parser_kwargs=parser_kwargs,
+            model_prompt=model_prompt,
+            model_id=model_id,
+            **batch_kwargs,
+        )
+
+    def _generate_logprobs(self, samples, **gen_kwargs):
+        prompts = [
+            self.model_prompt.format(prompt=self._format_prompt(sample))
+            for sample in samples
+        ]
+        generate_args = {**self.defaults, **gen_kwargs}
+
+        # verify if logprobs is mentioned in the generate_args, if not add it and return top10 logprobs
+        if "logprobs" not in generate_args:
+            generate_args["logprobs"] = 10
+
+        if self.server_supports_batched:
+            response = self.client.completions.create(prompt=prompts, **generate_args)
+            return [choice.logprobs.top_logprobs for choice in response.choices]
+
+        n = gen_kwargs.get("n", 1)
+        results = []
+        for prompt in prompts:
+            for _ in range(n):
+                response = self.client.completions.create(
+                    prompt=prompt, **generate_args
+                )
+                results.append(response.choices[0].logprobs.top_logprobs)
+        return results
+
+    def _parse(self, generations: List[List[Dict]]) -> List[List[str]]:
+        # override the parse method to convert the generations to json string
+        # convert the generations to json string to save as dataset
+        # this is because the dataset can only store key value pairs which are consistent
+        return [[json.dumps(item) for item in sublist] for sublist in generations]
+
+    def generate(self, samples: Dataset, **gen_kwargs) -> Dataset:
+        """
+        Generate the output from the block. This method should first validate the input data,
+        then generate the output, and finally parse the generated output before returning it.
+
+        :return: The parsed output after generation.
+        """
+        num_samples = self.block_config.get("num_samples", None)
+        logger.debug("Generating outputs for {} samples".format(len(samples)))
+
+        if (num_samples is not None) and ("num_samples" not in samples.column_names):
+            samples = samples.add_column("num_samples", [num_samples] * len(samples))
+
+        # validate each sample
+        # Log errors and remove invalid samples
+        valid_samples = []
+
+        for sample in samples:
+            if self._validate(self.prompt_template, sample):
+                valid_samples.append(sample)
+            else:
+                logger.warning(
+                    f"Sample failed validation: {sample}"
+                )  # Log details of the failed sample
+
+        samples = valid_samples
+
+        if len(samples) == 0:
+            logger.warning(
+                "No valid samples to generate outputs for, returning empty dataset"
+            )
+            return Dataset.from_list([])
+
+        # generate the output
+
+        outputs = self._generate_logprobs(samples, **gen_kwargs)
+        logger.debug("Generated outputs: %s", outputs)
+
+        output_dataset = Dataset.from_list(samples)
+        output_dataset = output_dataset.add_column(
+            self.output_cols[0],
+            self._parse(outputs),  # pylint: disable=no-value-for-parameter
+        )
+
+        return output_dataset
+
+
+@BlockRegistry.register("LLMMessagesBlock")
+class LLMMessagesBlock(Block):
+    def __init__(
+        self,
+        block_name,
+        client,
+        input_col,
+        output_col,
+        model_prompt=None, 
+        model_id=None,
+        **batch_kwargs,
+    ) -> None:
+        self.block_name = block_name
+        self.model_prompt = model_prompt
+        self.batch_params = batch_kwargs.get("batch_kwargs", {})
+        self.input_col = input_col
+        self.output_col = output_col
+        self.client = client
+
+        if model_id:
+            self.model = model_id
+        else:
+            self.model = self.client.models.list().data[0].id
+        
+        self.defaults = {
+            "model": self.model,
+            "temperature": 0,
+            "max_tokens": 4096,
+        }
+        self.server_supports_batched = server_supports_batched(client, self.model)
+
+    def _generate(self, samples, **gen_kwargs) -> list:
+        generate_args = {**self.defaults, **gen_kwargs}
+
+        if "n" in generate_args and generate_args.get("temperature", 0) <= 0:
+            generate_args["temperature"] = 0.7
+            logger.warning(
+                "Temperature should be greater than 0 for n > 1, setting temperature to 0.7"
+            )
+
+        messages = samples[self.input_col]
+
+        results = []
+        n = gen_kwargs.get("n", 1)
+        for message in messages:
+            responses = self.client.chat.completions.create(messages=message, **generate_args)
+            if n > 1:
+                results.append([choice.message.content for choice in responses.choices])
+            else:
+                results.append(responses.choices[0].message.content)
+        return results
+
+    def generate(self, samples: Dataset, **gen_kwargs) -> Dataset:
+        outputs = self._generate(samples, **gen_kwargs)
+        samples = samples.add_column(self.output_col, outputs)
+        return samples