dataset

dataset

DataFrame normalization and JSON schema inference for training data.

Provides utilities for standardizing DataFrames (type coercion, missing-value handling) and deriving JSON schemas used for validating generated records.

Functions:

Name	Description
check_enum_type	Return enum schema if the series is an enum, otherwise return None.
make_json_schema	Generate a JSON schema from the given DataFrame.
relax_numeric_bounds	Return a copy of schema with float (number) range bounds removed.
normalize_dataframe	Ensure DataFrame meets standards for use in Safe Synthesizer models.
normalize_column	Normalize the given pandas series.

check_enum_type(series, max_distinct=None, max_singletons=None)

Return enum schema if the series is an enum, otherwise return None.

Parameters:

Name	Type	Description	Default
series	Series	Data object to check for enum type.	required
max_distinct	int | float | None	Maximum number of distinct values to be considered an enum.	None
max_singletons	int | float | None	Maximum number of values with a single occurrence to be considered an enum.	None

Returns:

Type	Description
dict | None	The enum schema if the series is an enum, otherwise None.

Source code in src/nemo_safe_synthesizer/data_processing/dataset.py

def check_enum_type(
    series: pd.Series,
    max_distinct: int | float | None = None,
    max_singletons: int | float | None = None,
) -> dict | None:
    """Return enum schema if the series is an enum, otherwise return None.

    Args:
        series: Data object to check for enum type.
        max_distinct: Maximum number of distinct values to be considered an enum.
        max_singletons: Maximum number of values with a single occurrence to be
            considered an enum.

    Returns:
        The enum schema if the series is an enum, otherwise None.
    """
    if max_distinct is None:
        max_distinct = len(series) ** SCHEMA_ENUM_MAX_DISTINCT_EXP
    if max_singletons is None:
        max_singletons = len(series) ** SCHEMA_ENUM_MAX_SINGLETONS_EXP

    value_counts = series.value_counts(dropna=False)

    is_enum = value_counts.count() <= max_distinct and (value_counts == 1).sum() <= max_singletons

    return {"enum": [_handle_enum_value(v) for v in value_counts.index.sort_values()]} if is_enum else None

make_json_schema(df, string_length_multiple=STRING_LENGTH_MULTIPLE)

Generate a JSON schema from the given DataFrame.

Inspects each column to determine its JSON type, numeric range, string length bounds, or enum values. See https://json-schema.org for the schema specification.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame to derive a JSON schema from.	required
string_length_multiple	float	Multiplier applied to observed string lengths to set minLength (divided) and maxLength (multiplied) bounds in the schema.	STRING_LENGTH_MULTIPLE

Returns:

Type	Description
dict	A dictionary representing the JSON schema with type, properties,
dict	and required keys.

Source code in src/nemo_safe_synthesizer/data_processing/dataset.py

def make_json_schema(df: pd.DataFrame, string_length_multiple: float = STRING_LENGTH_MULTIPLE) -> dict:
    """Generate a JSON schema from the given DataFrame.

    Inspects each column to determine its JSON type, numeric range, string
    length bounds, or enum values. See https://json-schema.org for the
    schema specification.

    Args:
        df: DataFrame to derive a JSON schema from.
        string_length_multiple: Multiplier applied to observed string lengths
            to set ``minLength`` (divided) and ``maxLength`` (multiplied)
            bounds in the schema.

    Returns:
        A dictionary representing the JSON schema with ``type``, ``properties``,
        and ``required`` keys.
    """
    schema = {"type": "object", "properties": {}, "required": []}

    for col in df.columns:
        series = df[col]
        col_schema = check_enum_type(series) or {}

        if not col_schema:
            col_types = [JSON_TYPE_MAP.get(t, "string") for t in series.apply(lambda x: type(x).__name__).unique()]

            col_types = list(set(col_types))

            if series.isna().any():
                col_types.append("null")

            if set(col_types).issubset(["integer", "number"]):
                col_schema.update(
                    {
                        "type": col_types[0],  # actual element instead of list
                        "minimum": float(series.min()),
                        "maximum": float(series.max()),
                    }
                )
            elif col_types == ["string"]:
                str_length = series.astype(str).apply(len)
                col_schema.update(
                    {
                        "type": "string",
                        "minLength": round(str_length.min() / string_length_multiple),
                        "maxLength": round(string_length_multiple * str_length.max()),
                    }
                )
            else:
                col_schema.update({"type": col_types})

        schema["properties"][col] = col_schema

        if not series.isna().any():
            schema["required"].append(col)

    return schema

relax_numeric_bounds(schema)

Return a copy of schema with float (number) range bounds removed.

make_json_schema records each numeric column's exact observed minimum/maximum. Neither structured-generation backend can enforce a floating-point range (XGrammar and the regex builder bound integers and enums but emit an unconstrained token stream for number), so those bounds act purely as a post-generation validation gate. On wide float tables the per-field rejections compound and can reject nearly every record even though the values are otherwise well-formed. Dropping the number bounds turns that hard rejection into acceptance while leaving integer and enum constraints -- which the grammar does enforce -- untouched.

Only number-typed properties are relaxed; a property whose type list includes integer keeps its (grammar-enforced, whole-number) bounds so the XGrammar constraint stays valid.

Parameters:

Name	Type	Description	Default
schema	dict	A JSON schema as produced by make_json_schema.	required

Returns:

Type	Description
dict	A deep-ish copy of schema with minimum/maximum stripped from
dict	pure number properties. The input is not mutated.

Source code in src/nemo_safe_synthesizer/data_processing/dataset.py

def relax_numeric_bounds(schema: dict) -> dict:
    """Return a copy of ``schema`` with float (``number``) range bounds removed.

    ``make_json_schema`` records each numeric column's exact observed
    ``minimum``/``maximum``. Neither structured-generation backend can enforce a
    floating-point range (XGrammar and the regex builder bound integers and
    enums but emit an unconstrained token stream for ``number``), so those bounds
    act purely as a post-generation validation gate. On wide float tables the
    per-field rejections compound and can reject nearly every record even though
    the values are otherwise well-formed. Dropping the ``number`` bounds turns
    that hard rejection into acceptance while leaving integer and enum
    constraints -- which the grammar does enforce -- untouched.

    Only ``number``-typed properties are relaxed; a property whose type list
    includes ``integer`` keeps its (grammar-enforced, whole-number) bounds so the
    XGrammar constraint stays valid.

    Args:
        schema: A JSON schema as produced by ``make_json_schema``.

    Returns:
        A deep-ish copy of ``schema`` with ``minimum``/``maximum`` stripped from
        pure ``number`` properties. The input is not mutated.
    """
    relaxed = dict(schema)
    properties = relaxed.get("properties")
    if not isinstance(properties, dict):
        return relaxed

    new_properties: dict = {}
    for name, prop in properties.items():
        if isinstance(prop, dict):
            types = prop.get("type")
            types = types if isinstance(types, list) else [types]
            # Relax only pure floats; keep integer bounds (grammar-enforced and
            # required to be whole numbers) and enum constraints intact.
            if "number" in types and "integer" not in types:
                prop = {k: v for k, v in prop.items() if k not in ("minimum", "maximum")}
        new_properties[name] = prop
    relaxed["properties"] = new_properties
    return relaxed

normalize_dataframe(df)

Ensure DataFrame meets standards for use in Safe Synthesizer models.

Pandas may be used to construct a lot of odd DataFrames with weird corner cases that can violate assumptions of downstream code and other libraries. This includes differences creating a DataFrame from different sources (like csv vs json vs jsonl) or when a manually crafted DataFrame is provided, such as using SafeSynthesizerDataset for testing. Rather than be defensive in every model where we use SafeSynthesizerDataset, we do some standardization of all DataFrames here.

Enforced standards, i.e., assumptions models that use SafeSynthesizerDataset may make: - Every column has a single datatype, e.g. all float, all str, or all int, with the exception of missing values in object columns, where we keep the pandas behavior of representing missing with a float numpy.nan for now.

• Date, time, datetime, and timedelta types are converted to string for downstream consistency between tokenization and schema serialization. Decimal types are converted to float.

Source code in src/nemo_safe_synthesizer/data_processing/dataset.py

def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """Ensure DataFrame meets standards for use in Safe Synthesizer models.

    Pandas may be used to construct a lot of odd DataFrames with weird corner
    cases that can violate assumptions of downstream code and other libraries.
    This includes differences creating a DataFrame from different sources (like
    csv vs json vs jsonl) or when a manually crafted DataFrame is provided, such
    as using SafeSynthesizerDataset for testing. Rather than be defensive in every
     model where we use SafeSynthesizerDataset, we do some standardization of
    all DataFrames here.

    Enforced standards, i.e., assumptions models that use SafeSynthesizerDataset may make:
    - Every column has a single datatype, e.g. all float, all str, or all int,
      with the exception of missing values in object columns, where we keep the
      pandas behavior of representing missing with a float numpy.nan for now.

    - Date, time, datetime, and timedelta types are converted to string for
      downstream consistency between tokenization and schema serialization.
      Decimal types are converted to float.
    """
    column_series = {column_name: normalize_column(df[column_name]) for column_name in df.columns}
    return pd.DataFrame(column_series)

normalize_column(series)

Normalize the given pandas series.

Parameters:

Name	Type	Description	Default
series	Series	Series to normalize.	required

Returns:

Type	Description
Series	Normalized series.

Source code in src/nemo_safe_synthesizer/data_processing/dataset.py

def normalize_column(series: pd.Series) -> pd.Series:
    """Normalize the given pandas series.

    Args:
        series: Series to normalize.

    Returns:
        Normalized series.
    """
    series_type = pd.api.types.infer_dtype(series, skipna=True)
    if series_type in CONVERT_TO_STR_TYPES:
        return series.astype(str).mask(series.isna(), None)
    if series_type in CONVERT_TO_FLOAT_TYPES:
        return series.astype(float).mask(series.isna(), None)
    return series