sdg-hub 0.1.0a2.dev0__py3-none-any.whl

sdg_hub/flows/generation/skills/synth_skills.yaml

@@ -0,0 +1,59 @@
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_questions
+    config_path: configs/skills/freeform_questions.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - question
+    batch_kwargs:
+      num_samples: 30
+  drop_duplicates:
+    - question
+- block_type: LLMBlock
+  block_config:
+    block_name: eval_questions
+    config_path: configs/skills/evaluate_freeform_questions.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - evaluation
+      - score
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_questions
+    filter_column: score
+    filter_value: 1.0
+    operation: operator.eq
+    convert_dtype: float
+    batch_kwargs:
+      num_procs: 8
+  drop_columns:
+    - evaluation
+    - score
+    - num_samples
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_responses
+    config_path: configs/skills/freeform_responses.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - response
+- block_type: LLMBlock
+  block_config:
+    block_name: evaluate_qa_pair
+    config_path: configs/skills/evaluate_freeform_pair.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - evaluation
+      - score
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_qa_pair
+    filter_column: score
+    filter_value: 2.0
+    operation: operator.ge
+    convert_dtype: float
+    batch_kwargs:
+      num_procs: 8
+  drop_columns:
+    - evaluation
+    - score

sdg_hub/logger_config.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import os
+import logging
+
+# Third Party
+from rich.logging import RichHandler
+
+
+def setup_logger(name):
+    # Set up the logger
+    log_level = os.getenv("LOG_LEVEL", "INFO")
+    logging.basicConfig(
+        level=log_level,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[RichHandler()],
+    )
+    logger = logging.getLogger(name)
+    return logger

sdg_hub/pipeline.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# Third Party
+from datasets import Dataset
+
+# Local
+from .logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+class EmptyDatasetError(Exception):
+    pass
+
+
+class Pipeline:
+    def __init__(self, chained_blocks: list) -> None:
+        """
+        Initialize the Pipeline class with a configuration dictionary.
+        config_dict: the run config py or yaml loaded into a dictionary
+        """
+        # pipeline config is the run configuration that consists of the pipeline steps
+        self.chained_blocks = chained_blocks
+
+    def _drop_duplicates(self, dataset, cols):
+        """
+        Drop duplicates from the dataset based on the columns provided.
+        """
+        df = dataset.to_pandas()
+        df = df.drop_duplicates(subset=cols).reset_index(drop=True)
+        return Dataset.from_pandas(df)
+
+    def generate(self, dataset) -> Dataset:
+        """
+        Generate the dataset by running the pipeline steps.
+        dataset: the input dataset
+        """
+        for block_prop in self.chained_blocks:
+            block_type = block_prop["block_type"]
+            block_config = block_prop["block_config"]
+            drop_columns = block_prop.get("drop_columns", [])
+            gen_kwargs = block_prop.get("gen_kwargs", {})
+            drop_duplicates_cols = block_prop.get("drop_duplicates", False)
+            block = block_type(**block_config)
+
+            logger.debug("------------------------------------\n")
+            logger.debug("Running block: %s", block_config["block_name"])
+            logger.debug("Input dataset: %s", dataset)
+
+            dataset = block.generate(dataset, **gen_kwargs)
+
+            if len(dataset) == 0:
+                raise EmptyDatasetError(
+                    f"Pipeline stopped: Empty dataset after running block: {block_config['block_name']}"
+                )
+
+            drop_columns_in_ds = [e for e in drop_columns if e in dataset.column_names]
+            if drop_columns:
+                dataset = dataset.remove_columns(drop_columns_in_ds)
+
+            if drop_duplicates_cols:
+                dataset = self._drop_duplicates(dataset, cols=drop_duplicates_cols)
+
+            logger.debug("Output dataset: %s", dataset)
+            logger.debug("------------------------------------\n\n")
+
+        return dataset

sdg_hub/prompts.py

@@ -0,0 +1,17 @@
+# Local
+from .registry import PromptRegistry
+
+
+@PromptRegistry.register("blank")
+def blank_chat_template():
+    return """{{ messages }}"""
+
+
+@PromptRegistry.register("instructlab")
+def instructlab_chat_template():
+    return """{% for message in messages %}{% if message['role'] == 'pretraining' %}{{ '<|pretrain|>' + message['content'] + '<|endoftext|>' + '<|/pretrain|>' }}{% elif message['role'] == 'system' %}{{ '<|system|>' + '\n' + message['content'] + '\n' }}{% elif message['role'] == 'user' %}{{ '<|user|>' + '\n' + message['content'] + '\n' }}{% elif message['role'] == 'assistant' %}{{ '<|assistant|>' + '\n' + message['content'] + '<|endoftext|>' + ('' if loop.last else '\n') }}{% endif %}{% if loop.last and add_generation_prompt %}{{ '<|assistant|>' + '\n' }}{% endif %}{% endfor %}"""
+
+
+@PromptRegistry.register("mistralai")
+def mistral_chat_template():
+    return """{%- if messages[0]['role'] == 'system' %}\n    {%- set system_message = messages[0]['content'] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n\n<s>\n{%- for message in loop_messages %}\n    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}\n        {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}\n    {%- endif %}\n    {%- if message['role'] == 'user' %}\n        {%- if loop.first and system_message is defined %}\n            {{- ' [INST] ' + system_message + '\\n\\n' + message['content'] + ' [/INST]' }}\n        {%- else %}\n            {{- ' [INST] ' + message['content'] + ' [/INST]' }}\n        {%- endif %}\n    {%- elif message['role'] == 'assistant' %}\n        {{- ' ' + message['content'] + '</s>'}}\n    {%- else %}\n        {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}\n    {%- endif %}\n{%- endfor %}\n"""

sdg_hub/registry.py

@@ -0,0 +1,122 @@
+# Standard
+from typing import Union, List, Dict
+
+# Third Party
+from jinja2 import Template
+
+# Local
+from .logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+class BlockRegistry:
+    """Registry for block classes to avoid manual additions to block type map."""
+
+    _registry: Dict[str, type] = {}
+
+    @classmethod
+    def register(cls, block_name: str):
+        """
+        Decorator to register a block class under a specified name.
+
+        :param block_name: Name under which to register the block.
+        """
+
+        def decorator(block_class):
+            cls._registry[block_name] = block_class
+            logger.debug(
+                f"Registered block '{block_name}' with class '{block_class.__name__}'"
+            )
+            return block_class
+
+        return decorator
+
+    @classmethod
+    def get_registry(cls):
+        """
+        Retrieve the current registry map of block types.
+
+        :return: Dictionary of registered block names and classes.
+        """
+        logger.debug("Fetching the block registry map.")
+        return cls._registry
+
+
+class PromptRegistry:
+    """Registry for managing Jinja2 prompt templates."""
+
+    _registry: Dict[str, Template] = {}
+
+    @classmethod
+    def register(cls, name: str):
+        """Decorator to register a Jinja2 template function by name.
+
+        :param name: Name of the template to register.
+        :return: A decorator that registers the Jinja2 template function.
+        """
+
+        def decorator(func):
+            template_str = func()
+            cls._registry[name] = Template(template_str)
+            logger.debug(f"Registered prompt template '{name}'")
+            return func
+
+        return decorator
+
+    @classmethod
+    def get_template(cls, name: str) -> Template:
+        """Retrieve a Jinja2 template by name.
+
+        :param name: Name of the template to retrieve.
+        :return: The Jinja2 template instance.
+        """
+        if name not in cls._registry:
+            raise KeyError(f"Template '{name}' not found.")
+        logger.debug(f"Retrieving prompt template '{name}'")
+        return cls._registry[name]
+
+    @classmethod
+    def get_registry(cls):
+        """
+        Retrieve the current registry map of block types.
+
+        :return: Dictionary of registered block names and classes.
+        """
+        logger.debug("Fetching the block registry map.")
+        return cls._registry
+
+    @classmethod
+    def render_template(
+        cls,
+        name: str,
+        messages: Union[str, List[Dict[str, str]]],
+        add_generation_prompt: bool = True,
+    ) -> str:
+        """Render the template with the provided messages or query.
+
+        :param name: Name of the template to render.
+        :param messages: Either a single query string or a list of messages (each as a dict with 'role' and 'content').
+        :param add_generation_prompt: Whether to add a generation prompt at the end.
+        :return: The rendered prompt as a string.
+        """
+
+        # Special handling for "blank" template
+        if name == "blank":
+            if not isinstance(messages, str):
+                raise ValueError(
+                    "The 'blank' template can only be used with a single query string, not a list of messages."
+                )
+            return messages  # Return the query as-is without templating
+
+        # Get the template
+        template = cls.get_template(name)
+
+        # If `messages` is a string, wrap it in a list with a default user role
+        if isinstance(messages, str):
+            messages = [{"role": "user", "content": messages}]
+
+        # Render the template with the `messages` list
+        return template.render(
+            messages=messages, add_generation_prompt=add_generation_prompt
+        )

sdg_hub/sdg.py

@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List
+import traceback
+import uuid
+
+# Third Party
+from datasets import Dataset, load_dataset
+from datasets.data_files import EmptyDatasetError
+from tqdm import tqdm
+
+# Local
+from .logger_config import setup_logger
+from .pipeline import Pipeline
+from .utils.datautils import safe_concatenate_datasets
+
+
+logger = setup_logger(__name__)
+
+
+class SDG:
+    def __init__(
+        self, pipelines: List[Pipeline], num_workers=1, batch_size=None, save_freq=None
+    ) -> None:
+        self.pipelines = pipelines
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        self.save_freq = save_freq
+
+    def _split_dataset(self, dataset: Dataset, batch_size: int) -> List[Dataset]:
+        """Split the dataset into smaller batches."""
+        total_size = len(dataset)
+        num_batches = (total_size + batch_size - 1) // batch_size
+
+        batches = [
+            (i * batch_size, min((i + 1) * batch_size, total_size))
+            for i in tqdm(range(num_batches))
+        ]
+
+        return batches
+
+    def _get_missing_data(self, seed_data, generated_data):
+        # Get the common columns between the two datasets
+        common_columns = list(
+            set(seed_data.column_names) & set(generated_data.column_names)
+        )
+
+        # Extract the relevant data based on common columns
+        seed_data_common = seed_data.select_columns(common_columns)
+        generated_data_common = generated_data.select_columns(common_columns)
+
+        # Convert to Pandas DataFrames for easier comparison
+        seed_df = seed_data_common.to_pandas()
+        generated_df = generated_data_common.to_pandas()
+
+        # Identify missing rows
+        missing_df = seed_df[
+            ~seed_df.apply(tuple, 1).isin(generated_df.apply(tuple, 1))
+        ]
+
+        # Convert back to Dataset
+        missing_data = Dataset.from_pandas(missing_df, preserve_index=False)
+
+        return missing_data
+
+    def _save_intermediate_checkpoint(self, dataset, checkpoint_dir):
+        checkpoint_id = uuid.uuid4().hex
+        checkpoint_file = f"{checkpoint_dir}/data_checkpoint_{checkpoint_id}.jsonl"
+        logger.info(f"Saving checkpoint to {checkpoint_file}")
+        dataset.to_json(checkpoint_file, orient="records", lines=True)
+
+    @staticmethod
+    def _generate_data(pipelines, input_split, ds, i=None):
+        logger.info(f"Processing split {i}")
+        input_split = ds.select(range(input_split[0], input_split[1]))
+        try:
+            for pipeline in pipelines:
+                input_split = pipeline.generate(input_split)
+            return input_split
+        except Exception as e:
+            logger.error(f"Error processing split {i}: {e}")
+            traceback.print_exc()
+            return None
+
+    def generate(self, dataset: Dataset, checkpoint_dir=None) -> Dataset:
+        # check if checkpoint_dir exists
+        pre_generated_data = []
+        if checkpoint_dir is not None:
+            try:
+                # check if there are any existing checkpoints
+                pre_generated_data = load_dataset(
+                    "json", data_dir=checkpoint_dir, split="train"
+                )
+                logger.info(
+                    f"Loading existing checkpoints from {checkpoint_dir}, with {pre_generated_data.num_rows} rows"
+                )
+                seed_data = self._get_missing_data(dataset, pre_generated_data)
+                if seed_data.num_rows == 0:
+                    logger.info(
+                        f"All seed data has been generated, no missing rows found, returning data from {checkpoint_dir}"
+                    )
+                    return pre_generated_data
+                logger.info(f"Found {seed_data.num_rows} missing rows in the dataset")
+
+            except EmptyDatasetError:
+                logger.info(
+                    f"No existing checkpoints found in {checkpoint_dir}, generating from scratch"
+                )
+                seed_data = dataset
+
+        else:
+            seed_data = dataset
+
+        if not self.batch_size:
+            # If batch size is not provided, generate the dataset in a single pass
+            generated_dataset = seed_data
+            # generated_data is initialized with seed_data, and it gets updated with each pipeline
+            for pipeline in self.pipelines:
+                generated_dataset = pipeline.generate(seed_data)
+            return generated_dataset
+        
+        logger.info("Splitting the dataset into smaller batches")
+        input_splits = (
+            self._split_dataset(seed_data, self.batch_size)
+            if self.batch_size
+            else [seed_data]
+        )
+        logger.info(
+            f"Generating dataset with {len(input_splits)} splits, batch size {self.batch_size}, and {self.num_workers} workers"
+        )
+
+        generated_data = [pre_generated_data] if pre_generated_data else []
+        last_saved_split_index = 0  # To track the last saved split
+
+        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
+            futures = [
+                executor.submit(
+                    self._generate_data, self.pipelines, input_split, seed_data, i
+                )
+                for i, input_split in enumerate(input_splits)
+            ]
+
+            for i, future in enumerate(tqdm(as_completed(futures), total=len(futures))):
+                generated_data_split = future.result()  # Ensure each future completes
+
+                if generated_data_split:
+                    generated_data.append(generated_data_split)
+                    logger.info(f"Finished future processing split {i} \n\n")
+                    if self.save_freq and (i + 1) % self.save_freq == 0:
+                        # Save only the new splits since the last checkpoint
+                        new_splits = generated_data[last_saved_split_index : i + 1]
+                        checkpoint_dataset = safe_concatenate_datasets(new_splits)
+                        # check if checkpoint_dataset is not None
+                        if checkpoint_dataset:
+                            self._save_intermediate_checkpoint(
+                                checkpoint_dataset, checkpoint_dir
+                            )
+
+                            last_saved_split_index = i + 1
+
+        generated_dataset = safe_concatenate_datasets(generated_data)
+
+        return generated_dataset

sdg_hub/utils/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# This is part of the public API, and used by instructlab
+class GenerateException(Exception):
+    """An exception raised during generate step."""

sdg_hub/utils/chunking.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import List
+import logging
+import re
+
+# Third Party
+from langchain_text_splitters import Language, RecursiveCharacterTextSplitter
+
+_DEFAULT_CHUNK_OVERLAP = 100
+
+logger = logging.getLogger(__name__)
+
+
+def _num_tokens_from_words(num_words) -> int:
+    return int(num_words * 1.3)  # 1 word ~ 1.3 token
+
+
+def _num_chars_from_tokens(num_tokens) -> int:
+    return int(num_tokens * 4)  # 1 token ~ 4 English character
+
+
+def chunk_document(documents: List, server_ctx_size, chunk_word_count) -> List[str]:
+    """
+    Iterates over the documents and splits them into chunks based on the word count provided by the user.
+    Args:
+        documents (list): List of documents retrieved from git (can also consist of a single document).
+        server_ctx_size (int): Context window size of server.
+        chunk_word_count (int): Maximum number of words to chunk a document.
+    Returns:
+         List[str]: List of chunked documents.
+    """
+
+    # Checks for input type error
+    if isinstance(documents, str):
+        documents = [documents]
+
+    elif not isinstance(documents, list):
+        raise TypeError(
+            "Expected: documents to be a list, but got {}".format(type(documents))
+        )
+
+    no_tokens_per_doc = _num_tokens_from_words(chunk_word_count)
+    if no_tokens_per_doc > int(server_ctx_size - 1024):
+        raise ValueError(
+            "Error: {}".format(
+                str(
+                    f"Given word count ({chunk_word_count}) per doc will exceed the server context window size ({server_ctx_size})"
+                )
+            )
+        )
+    # Placeholder for params
+    content = []
+    chunk_size = _num_chars_from_tokens(no_tokens_per_doc)
+    chunk_overlap = _DEFAULT_CHUNK_OVERLAP
+
+    # Using Markdown as default, document-specific chunking will be implemented in seperate pr.
+    text_splitter = RecursiveCharacterTextSplitter.from_language(
+        language=Language.MARKDOWN,
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+    )
+
+    # Determine file type for heuristics, default with markdown
+    for docs in documents:
+        # Use regex to remove unnecessary dashes in front of pipe characters in a markdown table.
+        docs = re.sub(r"-{2,}\|", "-|", docs)
+        # Remove unnecessary spaces in front of pipe characters in a markdown table.
+        docs = re.sub(r"\  +\|", " |", docs)
+        temp = text_splitter.create_documents([docs])
+        content.extend([item.page_content for item in temp])
+    return content

sdg_hub/utils/datamixing.py

@@ -0,0 +1,123 @@
+# Standard
+import json
+import os
+
+# Third Party
+from datasets import Dataset, load_dataset
+import yaml
+
+# First Party
+from sdg_hub.logger_config import setup_logger
+from .datautils import safe_concatenate_datasets
+
+
+LOGGER = setup_logger(__name__)
+ALLOWED_COLS = ["id", "messages", "metadata"]
+
+
+def adjust_train_sample_size(ds: Dataset, num_samples: int):
+    LOGGER.info(f"Rebalancing dataset to have {num_samples} samples ...")
+    df = ds.to_pandas()
+    df = df.sample(n=num_samples, random_state=42, replace=True).reset_index(drop=True)
+    return Dataset.from_pandas(df)
+
+
+def load_ds(path, sampling_size):
+    if path.endswith(".jsonl"):
+        LOGGER.info(f"Loading dataset from {path} ...")
+        dataset = load_dataset("json", data_files=path, split="train")
+    else: 
+        LOGGER.info(f"Loading dataset from HF {path} ...")
+        dataset = load_dataset(path, split="train")
+    LOGGER.info(f"Dataset columns: {dataset.column_names}")
+    LOGGER.info(f"Dataset loaded with {len(dataset)} samples")
+
+    if sampling_size != 1.0:
+        if isinstance(sampling_size, int):
+            num_samples = sampling_size
+        else:
+            num_samples = int(len(dataset) * sampling_size)
+        dataset = adjust_train_sample_size(dataset, num_samples)
+
+    # move any column that is not in ALLOWED_COLS to metadata
+    def move_unallowed_cols_to_metadata(example):
+        metadata = example.get("metadata", {})
+        if isinstance(metadata, str):
+            metadata = json.loads(metadata)
+        for col in dataset.column_names:
+            if col not in ALLOWED_COLS:
+                metadata[col] = example[col]
+                example.pop(col)
+        example["metadata"] = json.dumps(metadata)
+        return example
+
+    dataset = dataset.map(move_unallowed_cols_to_metadata, num_proc=8)
+
+    # check if metadata column is string if not convert it using json.dumps
+    if not isinstance(dataset["metadata"][0], str):
+        dataset = dataset.map(
+            lambda x: {"metadata": json.dumps(x["metadata"])}, num_proc=8
+        )
+
+    return dataset
+
+
+def add_system_message(sample: dict, sys_prompt: str) -> dict:
+    # check if the messages have role system
+    has_system = False
+    for msg in sample["messages"]:
+        if msg["role"] == "system":
+            has_system = True
+            msg["content"] = sys_prompt
+
+    if not has_system:
+        sample["messages"].insert(0, {"role": "system", "content": sys_prompt})
+
+    return sample
+
+
+class Recipe:
+    def __init__(self, recipe_path):
+        self.recipe_path = recipe_path
+        self.recipe = self._load_recipe()
+        self.sys_prompt = self.recipe.get("sys_prompt", "")
+        self.dataset_added = False
+
+    def _load_recipe(self):
+        with open(self.recipe_path, encoding="utf-8") as fp:
+            return yaml.safe_load(fp)
+
+    def add_dataset(self, path, sampling_size=1.0):
+        self.dataset_added = True
+        self.recipe["datasets"].append({"path": path, "sampling_size": sampling_size})
+
+    def save_recipe(self, output_path):
+        # check if directory exists
+        output_dir = os.path.dirname(output_path)
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        with open(output_path, "w", encoding="utf-8") as fp:
+            yaml.dump(self.recipe, fp)
+
+    def save_mixed_dataset(self, output_path):
+        if not self.dataset_added:
+            LOGGER.error("No dataset added to the recipe")
+
+        mixed_ds = [
+            load_ds(dataset["path"], dataset["sampling_size"])
+            for dataset in self.recipe["datasets"]
+        ]
+
+        mixed_ds = safe_concatenate_datasets(mixed_ds)
+        mixed_ds = mixed_ds.map(
+            add_system_message, fn_kwargs={"sys_prompt": self.sys_prompt}, num_proc=8
+        )
+
+        # assert that the dataset only has the allowed columns
+        assert set(mixed_ds.column_names) == set(
+            ALLOWED_COLS
+        ), "Dataset has invalid columns"
+
+        mixed_ds.to_json(output_path, orient="records", lines=True)
+        LOGGER.info(f"Mixed Dataset saved to {output_path}")

sdg_hub/utils/datautils.py

@@ -0,0 +1,14 @@
+# Third Party
+from datasets import concatenate_datasets
+
+
+def safe_concatenate_datasets(datasets: list):
+    """
+    Concatenate datasets safely, ignoring any datasets that are None or empty.
+    """
+    filtered_datasets = [ds for ds in datasets if ds is not None and ds.num_rows > 0]
+
+    if not filtered_datasets:
+        return None
+
+    return concatenate_datasets(filtered_datasets)