data-designer 0.1.0__py3-none-any.whl

data_designer/config/seed.py

@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC
+from enum import Enum
+from typing import Optional, Union
+
+from pydantic import Field, field_validator, model_validator
+from typing_extensions import Self
+
+from .base import ConfigBase
+from .datastore import DatastoreSettings
+from .utils.io_helpers import (
+    VALID_DATASET_FILE_EXTENSIONS,
+    validate_dataset_file_path,
+    validate_path_contains_files_of_type,
+)
+
+
+class SamplingStrategy(str, Enum):
+    ORDERED = "ordered"
+    SHUFFLE = "shuffle"
+
+
+class IndexRange(ConfigBase):
+    start: int = Field(ge=0, description="The start index of the index range (inclusive)")
+    end: int = Field(ge=0, description="The end index of the index range (inclusive)")
+
+    @model_validator(mode="after")
+    def _validate_index_range(self) -> Self:
+        if self.start > self.end:
+            raise ValueError("'start' index must be less than or equal to 'end' index")
+        return self
+
+    @property
+    def size(self) -> int:
+        return self.end - self.start + 1
+
+
+class PartitionBlock(ConfigBase):
+    index: int = Field(default=0, ge=0, description="The index of the partition to sample from")
+    num_partitions: int = Field(default=1, ge=1, description="The total number of partitions in the dataset")
+
+    @model_validator(mode="after")
+    def _validate_partition_block(self) -> Self:
+        if self.index >= self.num_partitions:
+            raise ValueError("'index' must be less than 'num_partitions'")
+        return self
+
+    def to_index_range(self, dataset_size: int) -> IndexRange:
+        partition_size = dataset_size // self.num_partitions
+        start = self.index * partition_size
+
+        # For the last partition, extend to the end of the dataset to include remainder rows
+        if self.index == self.num_partitions - 1:
+            end = dataset_size - 1
+        else:
+            end = ((self.index + 1) * partition_size) - 1
+        return IndexRange(start=start, end=end)
+
+
+class SeedConfig(ConfigBase):
+    """Configuration for sampling data from a seed dataset.
+
+    Args:
+        dataset: Path or identifier for the seed dataset.
+        sampling_strategy: Strategy for how to sample rows from the dataset.
+            - ORDERED: Read rows sequentially in their original order.
+            - SHUFFLE: Randomly shuffle rows before sampling. When used with
+              selection_strategy, shuffling occurs within the selected range/partition.
+        selection_strategy: Optional strategy to select a subset of the dataset.
+            - IndexRange: Select a specific range of indices (e.g., rows 100-200).
+            - PartitionBlock: Select a partition by splitting the dataset into N equal parts.
+              Partition indices are zero-based (index=0 is the first partition, index=1 is
+              the second, etc.).
+
+    Examples:
+        Read rows sequentially from start to end:
+            SeedConfig(dataset="my_data.parquet", sampling_strategy=SamplingStrategy.ORDERED)
+
+        Read rows in random order:
+            SeedConfig(dataset="my_data.parquet", sampling_strategy=SamplingStrategy.SHUFFLE)
+
+        Read specific index range (rows 100-199):
+            SeedConfig(
+                dataset="my_data.parquet",
+                sampling_strategy=SamplingStrategy.ORDERED,
+                selection_strategy=IndexRange(start=100, end=199)
+            )
+
+        Read random rows from a specific index range (shuffles within rows 100-199):
+            SeedConfig(
+                dataset="my_data.parquet",
+                sampling_strategy=SamplingStrategy.SHUFFLE,
+                selection_strategy=IndexRange(start=100, end=199)
+            )
+
+        Read from partition 2 (3rd partition, zero-based) of 5 partitions (20% of dataset):
+            SeedConfig(
+                dataset="my_data.parquet",
+                sampling_strategy=SamplingStrategy.ORDERED,
+                selection_strategy=PartitionBlock(index=2, num_partitions=5)
+            )
+
+        Read shuffled rows from partition 0 of 10 partitions (shuffles within the partition):
+            SeedConfig(
+                dataset="my_data.parquet",
+                sampling_strategy=SamplingStrategy.SHUFFLE,
+                selection_strategy=PartitionBlock(index=0, num_partitions=10)
+            )
+    """
+
+    dataset: str
+    sampling_strategy: SamplingStrategy = SamplingStrategy.ORDERED
+    selection_strategy: Optional[Union[IndexRange, PartitionBlock]] = None
+
+
+class SeedDatasetReference(ABC, ConfigBase):
+    dataset: str
+
+
+class DatastoreSeedDatasetReference(SeedDatasetReference):
+    datastore_settings: DatastoreSettings
+
+    @property
+    def repo_id(self) -> str:
+        return "/".join(self.dataset.split("/")[:-1])
+
+    @property
+    def filename(self) -> str:
+        return self.dataset.split("/")[-1]
+
+
+class LocalSeedDatasetReference(SeedDatasetReference):
+    @field_validator("dataset", mode="after")
+    def validate_dataset_is_file(cls, v: str) -> str:
+        valid_wild_card_versions = {f"*{ext}" for ext in VALID_DATASET_FILE_EXTENSIONS}
+        if any(v.endswith(wildcard) for wildcard in valid_wild_card_versions):
+            parts = v.split("*.")
+            file_path = parts[0]
+            file_extension = parts[-1]
+            validate_path_contains_files_of_type(file_path, file_extension)
+        else:
+            validate_dataset_file_path(v)
+        return v

data_designer/config/utils/code_lang.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Union
+
+
+class CodeLang(str, Enum):
+    GO = "go"
+    JAVASCRIPT = "javascript"
+    JAVA = "java"
+    KOTLIN = "kotlin"
+    PYTHON = "python"
+    RUBY = "ruby"
+    RUST = "rust"
+    SCALA = "scala"
+    SWIFT = "swift"
+    TYPESCRIPT = "typescript"
+    SQL_SQLITE = "sql:sqlite"
+    SQL_TSQL = "sql:tsql"
+    SQL_BIGQUERY = "sql:bigquery"
+    SQL_MYSQL = "sql:mysql"
+    SQL_POSTGRES = "sql:postgres"
+    SQL_ANSI = "sql:ansi"
+
+    @staticmethod
+    def parse(value: Union[str, CodeLang]) -> tuple[str, Union[str, None]]:
+        value = value.value if isinstance(value, CodeLang) else value
+        split_vals = value.split(":")
+        return (split_vals[0], split_vals[1] if len(split_vals) > 1 else None)
+
+    @staticmethod
+    def parse_lang(value: Union[str, CodeLang]) -> str:
+        return CodeLang.parse(value)[0]
+
+    @staticmethod
+    def parse_dialect(value: Union[str, CodeLang]) -> Union[str, None]:
+        return CodeLang.parse(value)[1]
+
+    @staticmethod
+    def supported_values() -> set[str]:
+        return {lang.value for lang in CodeLang}
+
+
+SQL_DIALECTS: set[CodeLang] = {
+    CodeLang.SQL_SQLITE,
+    CodeLang.SQL_TSQL,
+    CodeLang.SQL_BIGQUERY,
+    CodeLang.SQL_MYSQL,
+    CodeLang.SQL_POSTGRES,
+    CodeLang.SQL_ANSI,
+}
+
+##########################################################
+# Helper functions
+##########################################################
+
+
+def code_lang_to_syntax_lexer(code_lang: Union[CodeLang, str]) -> str:
+    """Convert the code language to a syntax lexer for Pygments.
+
+    Reference: https://pygments.org/docs/lexers/
+    """
+    code_lang_to_lexer = {
+        CodeLang.GO: "golang",
+        CodeLang.JAVASCRIPT: "javascript",
+        CodeLang.JAVA: "java",
+        CodeLang.KOTLIN: "kotlin",
+        CodeLang.PYTHON: "python",
+        CodeLang.RUBY: "ruby",
+        CodeLang.RUST: "rust",
+        CodeLang.SCALA: "scala",
+        CodeLang.SWIFT: "swift",
+        CodeLang.SQL_SQLITE: "sql",
+        CodeLang.SQL_ANSI: "sql",
+        CodeLang.SQL_TSQL: "tsql",
+        CodeLang.SQL_BIGQUERY: "sql",
+        CodeLang.SQL_MYSQL: "mysql",
+        CodeLang.SQL_POSTGRES: "postgres",
+    }
+    return code_lang_to_lexer.get(code_lang, code_lang)

data_designer/config/utils/constants.py

@@ -0,0 +1,313 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from enum import Enum
+import os
+from pathlib import Path
+
+from rich.theme import Theme
+
+DEFAULT_NUM_RECORDS = 10
+
+EPSILON = 1e-8
+REPORTING_PRECISION = 2
+
+DEFAULT_REPR_HTML_STYLE = "nord"
+
+REPR_HTML_FIXED_WIDTH = 1000
+REPR_HTML_TEMPLATE = """
+<meta charset="UTF-8">
+<style>
+{{css}}
+
+.code {{{{
+  padding: 4px;
+  border: 1px solid grey;
+  border-radius: 4px;
+  max-width: {fixed_width}px;
+  width: 100%;
+  display: inline-block;
+  box-sizing: border-box;
+  text-align: left;
+  vertical-align: top;
+  line-height: normal;
+  overflow-x: auto;
+}}}}
+
+.code pre {{{{
+  white-space: pre-wrap;       /* CSS 3 */
+  white-space: -moz-pre-wrap;  /* Mozilla, since 1999 */
+  white-space: -pre-wrap;      /* Opera 4-6 */
+  white-space: -o-pre-wrap;    /* Opera 7 */
+  word-wrap: break-word;
+  overflow-wrap: break-word;
+  margin: 0;
+}}}}
+</style>
+{{highlighted_html}}
+""".format(fixed_width=REPR_HTML_FIXED_WIDTH)
+
+
+class NordColor(Enum):
+    NORD0 = "#2E3440"  # Darkest gray (background)
+    NORD1 = "#3B4252"  # Dark gray
+    NORD2 = "#434C5E"  # Medium dark gray
+    NORD3 = "#4C566A"  # Lighter dark gray
+    NORD4 = "#D8DEE9"  # Light gray (default text)
+    NORD5 = "#E5E9F0"  # Very light gray
+    NORD6 = "#ECEFF4"  # Almost white
+    NORD7 = "#8FBCBB"  # Teal
+    NORD8 = "#88C0D0"  # Light cyan
+    NORD9 = "#81A1C1"  # Soft blue
+    NORD10 = "#5E81AC"  # Darker blue
+    NORD11 = "#BF616A"  # Red
+    NORD12 = "#D08770"  # Orange
+    NORD13 = "#EBCB8B"  # Yellow
+    NORD14 = "#A3BE8C"  # Green
+    NORD15 = "#B48EAD"  # Purple
+
+
+RICH_CONSOLE_THEME = Theme(
+    {
+        "repr.number": NordColor.NORD15.value,  # Purple for numbers
+        "repr.string": NordColor.NORD14.value,  # Green for strings
+        "repr.bool_true": NordColor.NORD9.value,  # Blue for True
+        "repr.bool_false": NordColor.NORD9.value,  # Blue for False
+        "repr.none": NordColor.NORD11.value,  # Red for None
+        "repr.brace": NordColor.NORD7.value,  # Teal for brackets/braces
+        "repr.comma": NordColor.NORD7.value,  # Teal for commas
+        "repr.ellipsis": NordColor.NORD7.value,  # Teal for ellipsis
+        "repr.attrib_name": NordColor.NORD3.value,  # Light gray for dict keys
+        "repr.attrib_equal": NordColor.NORD7.value,  # Teal for equals signs
+        "repr.call": NordColor.NORD10.value,  # Darker blue for function calls
+        "repr.function_name": NordColor.NORD10.value,  # Darker blue for function names
+        "repr.class_name": NordColor.NORD12.value,  # Orange for class names
+        "repr.module_name": NordColor.NORD8.value,  # Light cyan for module names
+        "repr.error": NordColor.NORD11.value,  # Red for errors
+        "repr.warning": NordColor.NORD13.value,  # Yellow for warnings
+    }
+)
+
+DEFAULT_HIST_NAME_COLOR = "medium_purple1"
+
+DEFAULT_HIST_VALUE_COLOR = "pale_green3"
+
+
+DEFAULT_AGE_RANGE = [18, 114]
+MIN_AGE = 0
+MAX_AGE = 114
+
+LOCALES_WITH_MANAGED_DATASETS = ["en_US", "ja_JP", "en_IN", "hi_IN"]
+
+US_STATES_AND_MAJOR_TERRITORIES = {
+    # States
+    "AK",
+    "AL",
+    "AR",
+    "AZ",
+    "CA",
+    "CO",
+    "CT",
+    "DE",
+    "FL",
+    "GA",
+    "HI",
+    "IA",
+    "ID",
+    "IL",
+    "IN",
+    "KS",
+    "KY",
+    "LA",
+    "MA",
+    "MD",
+    "ME",
+    "MI",
+    "MN",
+    "MO",
+    "MS",
+    "MT",
+    "NC",
+    "ND",
+    "NE",
+    "NH",
+    "NJ",
+    "NM",
+    "NV",
+    "NY",
+    "OH",
+    "OK",
+    "OR",
+    "PA",
+    "RI",
+    "SC",
+    "SD",
+    "TN",
+    "TX",
+    "UT",
+    "VA",
+    "VT",
+    "WA",
+    "WI",
+    "WV",
+    "WY",
+    # D.C.
+    "DC",
+    # Territories
+    "AS",
+    "GU",
+    "MP",
+    "PR",
+    "VI",
+}
+
+MAX_TEMPERATURE = 2.0
+MIN_TEMPERATURE = 0.0
+MAX_TOP_P = 1.0
+MIN_TOP_P = 0.0
+MIN_MAX_TOKENS = 1
+REASONING_TRACE_COLUMN_POSTFIX = "__reasoning_trace"
+
+AVAILABLE_LOCALES = [
+    "ar_AA",
+    "ar_AE",
+    "ar_BH",
+    "ar_EG",
+    "ar_JO",
+    "ar_PS",
+    "ar_SA",
+    "az_AZ",
+    "bg_BG",
+    "bn_BD",
+    "bs_BA",
+    "cs_CZ",
+    "da_DK",
+    "de",
+    "de_AT",
+    "de_CH",
+    "de_DE",
+    "dk_DK",
+    "el_CY",
+    "el_GR",
+    "en",
+    "en_AU",
+    "en_BD",
+    "en_CA",
+    "en_GB",
+    "en_IE",
+    "en_IN",
+    "en_NZ",
+    "en_PH",
+    "en_TH",
+    "en_US",
+    "es",
+    "es_AR",
+    "es_CA",
+    "es_CL",
+    "es_CO",
+    "es_ES",
+    "es_MX",
+    "et_EE",
+    "fa_IR",
+    "fi_FI",
+    "fil_PH",
+    "fr_BE",
+    "fr_CA",
+    "fr_CH",
+    "fr_FR",
+    #    "fr_QC", deprecated, use fr_CA instead
+    "ga_IE",
+    "he_IL",
+    "hi_IN",
+    "hr_HR",
+    "hu_HU",
+    "hy_AM",
+    "id_ID",
+    "it_CH",
+    "it_IT",
+    "ja_JP",
+    "ka_GE",
+    "ko_KR",
+    "la",
+    "lb_LU",
+    "lt_LT",
+    "lv_LV",
+    "mt_MT",
+    "ne_NP",
+    "nl_BE",
+    "nl_NL",
+    "no_NO",
+    "or_IN",
+    "pl_PL",
+    "pt_BR",
+    "pt_PT",
+    "ro_RO",
+    "ru_RU",
+    "sk_SK",
+    "sl_SI",
+    "sq_AL",
+    "sv_SE",
+    "ta_IN",
+    "th",
+    "th_TH",
+    "tl_PH",
+    "tr_TR",
+    "tw_GH",
+    "uk_UA",
+    "vi_VN",
+    "zh_CN",
+    "zh_TW",
+    "zu_ZA",
+]
+
+DATA_DESIGNER_HOME_ENV_VAR = "DATA_DESIGNER_HOME"
+
+DATA_DESIGNER_HOME = Path(os.getenv(DATA_DESIGNER_HOME_ENV_VAR, Path.home() / ".data-designer"))
+
+MANAGED_ASSETS_PATH_ENV_VAR = "DATA_DESIGNER_MANAGED_ASSETS_PATH"
+
+MANAGED_ASSETS_PATH = Path(os.getenv(MANAGED_ASSETS_PATH_ENV_VAR, DATA_DESIGNER_HOME / "managed-assets"))
+
+MODEL_CONFIGS_FILE_NAME = "model_configs.yaml"
+
+MODEL_CONFIGS_FILE_PATH = DATA_DESIGNER_HOME / MODEL_CONFIGS_FILE_NAME
+
+MODEL_PROVIDERS_FILE_NAME = "model_providers.yaml"
+
+MODEL_PROVIDERS_FILE_PATH = DATA_DESIGNER_HOME / MODEL_PROVIDERS_FILE_NAME
+
+NVIDIA_PROVIDER_NAME = "nvidia"
+
+NVIDIA_API_KEY_ENV_VAR_NAME = "NVIDIA_API_KEY"
+
+OPENAI_PROVIDER_NAME = "openai"
+
+OPENAI_API_KEY_ENV_VAR_NAME = "OPENAI_API_KEY"
+
+PREDEFINED_PROVIDERS = [
+    {
+        "name": NVIDIA_PROVIDER_NAME,
+        "endpoint": "https://integrate.api.nvidia.com/v1",
+        "provider_type": "openai",
+        "api_key": NVIDIA_API_KEY_ENV_VAR_NAME,
+    },
+    {
+        "name": OPENAI_PROVIDER_NAME,
+        "endpoint": "https://api.openai.com/v1",
+        "provider_type": "openai",
+        "api_key": OPENAI_API_KEY_ENV_VAR_NAME,
+    },
+]
+
+PREDEFINED_PROVIDERS_MODEL_MAP = {
+    NVIDIA_PROVIDER_NAME: {
+        "text": "nvidia/nvidia-nemotron-nano-9b-v2",
+        "reasoning": "openai/gpt-oss-20b",
+        "vision": "nvidia/nemotron-nano-12b-v2-vl",
+    },
+    OPENAI_PROVIDER_NAME: {
+        "text": "gpt-4.1",
+        "reasoning": "gpt-5",
+        "vision": "gpt-5",
+    },
+}

data_designer/config/utils/errors.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from ...errors import DataDesignerError
+
+
+class UserJinjaTemplateSyntaxError(DataDesignerError): ...
+
+
+class InvalidEnumValueError(DataDesignerError): ...
+
+
+class InvalidTypeUnionError(DataDesignerError): ...
+
+
+class InvalidDiscriminatorFieldError(DataDesignerError): ...
+
+
+class DatasetSampleDisplayError(DataDesignerError): ...

data_designer/config/utils/info.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Literal, TypeVar
+
+from ..models import ModelConfig, ModelProvider
+from ..sampler_params import SamplerType
+from .type_helpers import get_sampler_params
+from .visualization import display_model_configs_table, display_model_providers_table, display_sampler_table
+
+
+class InfoType(str, Enum):
+    SAMPLERS = "samplers"
+    MODEL_CONFIGS = "model_configs"
+    MODEL_PROVIDERS = "model_providers"
+
+
+ConfigBuilderInfoType = Literal[InfoType.SAMPLERS, InfoType.MODEL_CONFIGS]
+DataDesignerInfoType = Literal[InfoType.MODEL_PROVIDERS]
+InfoTypeT = TypeVar("InfoTypeT", bound=InfoType)
+
+
+class InfoDisplay(ABC):
+    """Base class for info display classes that provide type-safe display methods."""
+
+    @abstractmethod
+    def display(self, info_type: InfoTypeT, **kwargs) -> None:
+        """Display information based on the provided info type.
+
+        Args:
+            info_type: Type of information to display.
+        """
+        ...
+
+
+class ConfigBuilderInfo(InfoDisplay):
+    def __init__(self, model_configs: list[ModelConfig]):
+        self._sampler_params = get_sampler_params()
+        self._model_configs = model_configs
+
+    def display(self, info_type: ConfigBuilderInfoType, **kwargs) -> None:
+        """Display information based on the provided info type.
+
+        Args:
+            info_type: Type of information to display. Only SAMPLERS and MODEL_CONFIGS are supported.
+
+        Raises:
+            ValueError: If an unsupported info_type is provided.
+        """
+        if info_type == InfoType.SAMPLERS:
+            self._display_sampler_info(sampler_type=kwargs.get("sampler_type"))
+        elif info_type == InfoType.MODEL_CONFIGS:
+            display_model_configs_table(self._model_configs)
+        else:
+            raise ValueError(
+                f"Unsupported info_type: {str(info_type)!r}. "
+                f"ConfigBuilderInfo only supports {InfoType.SAMPLERS.value!r} and {InfoType.MODEL_CONFIGS.value!r}."
+            )
+
+    def _display_sampler_info(self, sampler_type: SamplerType | None) -> None:
+        if sampler_type is not None:
+            title = f"{SamplerType(sampler_type).value.replace('_', ' ').title()} Sampler"
+            display_sampler_table({sampler_type: self._sampler_params[sampler_type]}, title=title)
+        else:
+            display_sampler_table(self._sampler_params)
+
+
+class InterfaceInfo(InfoDisplay):
+    def __init__(self, model_providers: list[ModelProvider]):
+        self._model_providers = model_providers
+
+    def display(self, info_type: DataDesignerInfoType, **kwargs) -> None:
+        """Display information based on the provided info type.
+
+        Args:
+            info_type: Type of information to display. Only MODEL_PROVIDERS is supported.
+
+        Raises:
+            ValueError: If an unsupported info_type is provided.
+        """
+        if info_type == InfoType.MODEL_PROVIDERS:
+            display_model_providers_table(self._model_providers)
+        else:
+            raise ValueError(
+                f"Unsupported info_type: {str(info_type)!r}. InterfaceInfo only supports {InfoType.MODEL_PROVIDERS.value!r}."
+            )