TimeGan Synthesizer - 'dict' object has no attribute 'dataframe'

[atulkakrana]

Description

Running bundled tutorial tutorial4_time_series.ipynb raises an error at the synthesis step.

How to Reproduce

Run notebook tutorial4_time_series.ipynb. This step gives an error: syn_model.generate(count=100).dataframe()

System Information

• OS: Ubuntu
• OS Version: 24.04 LTS
• Language Version: Python 3.10, Torch 2.10

Additional Context

Torch 2.10 is not recommended but using torch version < 2.4 the gives an import error related torch nn missing.

Error

AttributeError                            Traceback (most recent call last)
Cell In[7], [line 1](vscode-notebook-cell:?execution_count=7&line=1)
----> [1](vscode-notebook-cell:?execution_count=7&line=1) syn_model.generate(count=100).dataframe()

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:56, in validate_arguments.<locals>.validate.<locals>.wrapper_function(*args, **kwargs)
     54 @wraps(_func)
     55 def wrapper_function(*args: Any, **kwargs: Any) -> Any:
---> [56](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:56)     return vd.call(*args, **kwargs)

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:151, in ValidatedFunction.call(self, *args, **kwargs)
    149 def call(self, *args: Any, **kwargs: Any) -> Any:
    150     m = self.init_model_instance(*args, **kwargs)
--> [151](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:151)     return self.execute(m)

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:227, in ValidatedFunction.execute(self, m)
    225     return self.raw_function(*args_, **kwargs, **var_kwargs)
    226 else:
--> [227](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:227)     return self.raw_function(**d, **var_kwargs)

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/synthcity/plugins/core/plugin.py:349, in Plugin.generate(self, count, constraints, random_state, **kwargs)
    345     gen_constraints = gen_constraints.extend(constraints)
    347 syn_schema = Schema.from_constraints(gen_constraints)
--> [349](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/synthcity/plugins/core/plugin.py:349) X_syn = self._generate(count=count, syn_schema=syn_schema, **kwargs)
    351 if X_syn.is_tabular():
    352     if self.compress_dataset:

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/synthcity/plugins/time_series/plugin_timegan.py:437, in TimeGANPlugin._generate(self, count, syn_schema, **kwargs)
    434     else:
    435         raise RuntimeError("unknow data type")
--> [437](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/synthcity/plugins/time_series/plugin_timegan.py:437) return self._safe_generate_time_series(_sample, count, syn_schema)

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:56, in validate_arguments.<locals>.validate.<locals>.wrapper_function(*args, **kwargs)
     54 @wraps(_func)
     55 def wrapper_function(*args: Any, **kwargs: Any) -> Any:
---> [56](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:56)     return vd.call(*args, **kwargs)

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:150, in ValidatedFunction.call(self, *args, **kwargs)
    149 def call(self, *args: Any, **kwargs: Any) -> Any:
--> [150](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:150)     m = self.init_model_instance(*args, **kwargs)
    151     return self.execute(m)

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:147, in ValidatedFunction.init_model_instance(self, *args, **kwargs)
    145 def init_model_instance(self, *args: Any, **kwargs: Any) -> BaseModel:
    146     values = self.build_values(args, kwargs)
--> [147](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/pydantic/deprecated/decorator.py:147)     return self.model(**values)

    [... skipping hidden 1 frame]

File ~/miniconda3/envs/syncity/lib/python3.10/site-packages/synthcity/plugins/core/schema.py:71, in Schema.initialize_domain(cls, model)
     68 @model_validator(mode="after")
     69 def initialize_domain(cls, model: "Schema") -> "Schema":
     70     if model.data is not None:
---> [71](https://vscode-remote+ssh-002dremote-002bneau-002dgpu-002d1.vscode-resource.vscode-cdn.net/home/atul/git/synthcity/tutorials/~/miniconda3/envs/syncity/lib/python3.10/site-packages/synthcity/plugins/core/schema.py:71)         X = model.data.dataframe()
     72         model.domain = model._infer_domain(
     73             X,
     74             sampling_strategy=model.sampling_strategy,
     75             random_state=model.random_state,
     76         )
     77         # Remove 'data' attribute from the model

AttributeError: 'dict' object has no attribute 'dataframe'

[M-Zeeb]

I had the same issue. It seems to be a problem with the pydantic version.
Fix for me, (using pydantic 2.12.2) was to declare @classmethod in a new line after all @field_validator and @model_validator in the files schema.py and distribution.py (under plugins/core/).

[atulkakrana]

Thanks @M-Zeeb. Can you paste your code here?
I will make the same updates on my side to test if that fixes the problem.

[M-Zeeb]

Sure, here you go:

distribution.py

# stdlib
from abc import ABCMeta, abstractmethod
from datetime import datetime, timedelta, timezone
from typing import Any, List, Optional, Tuple

# third party
import numpy as np
import pandas as pd
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    FieldValidationInfo,
    PrivateAttr,
    ValidationInfo,
    field_validator,
    model_validator,
)

# synthcity absolute
from synthcity.plugins.core.constraints import Constraints

Rule = Tuple[str, str, Any]  # Define a type alias for clarity


class Distribution(BaseModel, metaclass=ABCMeta):
    """
    .. inheritance-diagram:: synthcity.plugins.core.distribution.Distribution
        :parts: 1

    Base class of all Distributions.

    The Distribution class characterizes the **empirical** marginal distribution of the feature.
    Each derived class must implement the following methods:
        get() - Return the metadata of the Distribution.
        sample() - Sample a value from the Distribution.
        includes() - Test if another Distribution is included in the local one.
        has() - Test if a value is included in the support of the Distribution.
        as_constraint() - Convert the Distribution to a set of Constraints.
        min() - Return the minimum of the support.
        max() - Return the maximum of the support.
        __eq__() - Testing equality of two Distributions.
        dtype() - Return the data type

    Examples of derived classes include CategoricalDistribution, FloatDistribution, and IntegerDistribution.
    """

    name: str
    data: Optional[pd.Series] = None
    random_state: Optional[int] = None
    sampling_strategy: str = "marginal"
    _rng: np.random.Generator = PrivateAttr()
    # DP parameters
    marginal_distribution: Optional[pd.Series] = None

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @field_validator("marginal_distribution", mode="before")
    @classmethod
    def _validate_marginal_distribution(
        cls: Any, v: Any, values: FieldValidationInfo
    ) -> Optional[pd.Series]:
        if "data" not in values.data or values.data["data"] is None:
            return v

        data = values.data["data"]
        if not isinstance(data, pd.Series):
            raise ValueError(f"Invalid data type {type(data)}")

        marginal = data.value_counts(dropna=False)
        del values["data"]

        return marginal

    @model_validator(mode="after")
    @classmethod
    def initialize_rng(cls, model: "Distribution") -> "Distribution":
        """
        Initializes the random number generator after model validation.
        """
        if model.random_state is not None:
            model._rng = np.random.default_rng(model.random_state)
        else:
            model._rng = np.random.default_rng()
        return model

    def marginal_states(self) -> Optional[List]:
        if self.marginal_distribution is None:
            return None

        return self.marginal_distribution.index.values

    def marginal_probabilities(self) -> Optional[List]:
        if self.marginal_distribution is None:
            return None

        return (
            self.marginal_distribution.values / self.marginal_distribution.values.sum()
        )

    def sample_marginal(self, count: int = 1) -> Any:
        if self.marginal_distribution is None:
            return None

        return self._rng.choice(
            self.marginal_states(),
            count,
            p=self.marginal_probabilities(),
        ).tolist()

    @abstractmethod
    def get(self) -> List[Any]:
        """Return the metadata of the Distribution."""
        ...

    @abstractmethod
    def sample(self, count: int = 1) -> Any:
        """Sample a value from the Distribution."""
        ...

    @abstractmethod
    def includes(self, other: "Distribution") -> bool:
        """Test if another Distribution is included in the local one."""
        ...

    @abstractmethod
    def has(self, val: Any) -> bool:
        """Test if a value is included in the Distribution."""
        ...

    @abstractmethod
    def as_constraint(self) -> Constraints:
        """Convert the Distribution to a set of Constraints."""
        ...

    @abstractmethod
    def min(self) -> Any:
        """Get the min value of the distribution."""
        ...

    @abstractmethod
    def max(self) -> Any:
        """Get the max value of the distribution."""
        ...

    def __eq__(self, other: Any) -> bool:
        return type(self) == type(other) and self.get() == other.get()

    def __contains__(self, item: Any) -> bool:
        """
        Example:
        >>> dist = CategoricalDistribution(name="foo", choices=["a", "b", "c"])
        >>> "a" in dist
        True
        """
        return self.has(item)

    @abstractmethod
    def dtype(self) -> str:
        ...


class CategoricalDistribution(Distribution):
    """
    .. inheritance-diagram:: synthcity.plugins.core.distribution.CategoricalDistribution
        :parts: 1
    """

    data: Optional[pd.Series] = None
    marginal_distribution: Optional[pd.Series] = None
    choices: List[Any] = Field(default_factory=list)

    @model_validator(mode="after")
    @classmethod
    def validate_and_initialize(
        cls, model: "CategoricalDistribution"
    ) -> "CategoricalDistribution":
        """
        Validates and initializes choices and marginal_distribution based on data or provided choices.
        Ensures that choices are unique and sorted.
        """
        if model.data is not None:
            # Set marginal_distribution based on data
            model.marginal_distribution = model.data.value_counts(normalize=True)
            model.choices = model.marginal_distribution.index.tolist()
        elif model.choices is not None:
            # Ensure choices are unique and sorted
            model.choices = sorted(set(model.choices))
            # Set uniform probabilities
            probabilities = np.ones(len(model.choices)) / len(model.choices)
            model.marginal_distribution = pd.Series(probabilities, index=model.choices)
        else:
            raise ValueError(
                "Invalid CategoricalDistribution: Provide either 'data' or 'choices'."
            )

        # Additional validation to ensure consistency
        if not isinstance(model.choices, list) or len(model.choices) == 0:
            raise ValueError(
                "CategoricalDistribution must have a non-empty 'choices' list."
            )
        if not isinstance(model.marginal_distribution, pd.Series):
            raise ValueError(
                "CategoricalDistribution must have a valid 'marginal_distribution'."
            )
        if len(model.choices) != len(model.marginal_distribution):
            raise ValueError(
                "'choices' and 'marginal_distribution' must have the same length."
            )

        return model

    def sample(self, count: int = 1) -> Any:
        """
        Samples values from the distribution based on the specified sampling strategy.
        If the distribution has only one choice, returns an array filled with that value.
        """
        if self.choices is not None and len(self.choices) == 1:
            samples = np.full(count, self.choices[0])
        else:
            if self.sampling_strategy == "marginal":
                if self.marginal_distribution is None:
                    raise ValueError(
                        "Cannot sample based on marginal distribution: marginal_distribution is not provided."
                    )
                return self._rng.choice(
                    self.marginal_distribution.index,
                    size=count,
                    p=self.marginal_distribution.values,
                )
            elif self.sampling_strategy == "uniform":
                return self._rng.choice(self.choices, size=count)
            else:
                raise ValueError(
                    f"Unsupported sampling strategy '{self.sampling_strategy}'."
                )
        return samples

    def get(self) -> List[Any]:
        """
        Returns the metadata of the distribution.
        """
        return [self.name, self.choices]

    def has(self, val: Any) -> bool:
        """
        Checks if a value is among the distribution's choices.
        """
        return val in self.choices

    def includes(self, other: "Distribution") -> bool:
        """
        Checks if another categorical distribution's choices are a subset of this distribution's choices.
        """
        if not isinstance(other, CategoricalDistribution):
            return False
        return set(other.choices).issubset(set(self.choices))

    def as_constraint(self) -> Constraints:
        """
        Converts the distribution to a set of constraints.
        """
        return Constraints(rules=[(self.name, "in", list(self.choices))])

    def min(self) -> Any:
        """
        Returns the minimum value among the choices.
        """
        return min(self.choices)

    def max(self) -> Any:
        """
        Returns the maximum value among the choices.
        """
        return max(self.choices)

    def dtype(self) -> str:
        """
        Determines the data type based on the choices.
        """
        types = {
            "object": 0,
            "float": 0,
            "int": 0,
        }
        for v in self.choices:
            if isinstance(v, float):
                types["float"] += 1
            elif isinstance(v, int):
                types["int"] += 1
            else:
                types["object"] += 1

        for t in types:
            if types[t] != 0:
                return t

        return "object"


class FloatDistribution(Distribution):
    """
    .. inheritance-diagram:: synthcity.plugins.core.distribution.FloatDistribution
        :parts: 1
    """

    low: Optional[float] = Field(default=None)
    high: Optional[float] = Field(default=None)
    _is_constant: bool = PrivateAttr(False)

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @model_validator(mode="after")
    @classmethod
    def validate_and_initialize(cls, model: "FloatDistribution") -> "FloatDistribution":
        """
        Validates and initializes the distribution.
        Sets '_is_constant' based on whether 'low' equals 'high'.
        Initializes 'marginal_distribution' based on 'data' if provided.
        """
        if model.data is not None:
            # Initialize marginal_distribution based on data
            # For float data, use value_counts(normalize=True) if data has repeated values
            # This will create a discrete approximation of the distribution
            model.marginal_distribution = model.data.value_counts(
                normalize=True
            ).sort_index()
            model.low = float(model.data.min())
            model.high = float(model.data.max())
        elif model.marginal_distribution is not None:
            # Set 'low' and 'high' based on marginal_distribution
            model.low = float(model.marginal_distribution.index.min())
            model.high = float(model.marginal_distribution.index.max())
        else:
            # Ensure 'low' and 'high' are provided
            if model.low is None or model.high is None:
                raise ValueError(
                    "FloatDistribution requires 'low' and 'high' values if 'data' or 'marginal_distribution' is not provided."
                )

        # Validate that low <= high
        if model.low > model.high:
            raise ValueError(
                f"Invalid range for '{model.name}': low ({model.low}) cannot be greater than high ({model.high})."
            )

        # Set _is_constant based on low == high
        model._is_constant = model.low == model.high

        # Ensure that low and high are finite numbers
        if not np.isfinite(model.low) or not np.isfinite(model.high):
            raise ValueError(
                f"Invalid range for '{model.name}': low or high is not finite (low={model.low}, high={model.high})."
            )

        return model

    def sample(self, count: int = 1) -> Any:
        """
        Samples values from the distribution.
        If the distribution is constant, returns an array filled with the constant value.
        Otherwise, samples based on the marginal distribution or uniform sampling.
        """
        if self._is_constant:
            if self.low is None:
                raise ValueError(
                    "Cannot sample: 'low' is None for a constant distribution."
                )
            samples = np.full(count, self.low)
        else:
            if self.low is None or self.high is None:
                raise ValueError("Cannot sample: 'low' or 'high' is None.")
            if (
                self.sampling_strategy == "marginal"
                and self.marginal_distribution is not None
            ):
                # Sample based on marginal distribution
                return self._rng.choice(
                    self.marginal_distribution.index.values,
                    size=count,
                    p=self.marginal_distribution.values,
                )
            else:
                # Proceed with uniform sampling
                samples = self._rng.uniform(low=self.low, high=self.high, size=count)
        return samples

    def get(self) -> List[Any]:
        """
        Returns the metadata of the distribution.
        """
        return [self.name, self.low, self.high]

    def has(self, val: Any) -> bool:
        """
        Checks if a value is within the distribution's range.
        """
        return self.low <= val <= self.high

    def includes(self, other: "Distribution") -> bool:
        """
        Checks if another distribution is entirely within this distribution.
        """
        if self.min() is None or self.max() is None:
            return False
        if other.min() is None or other.max() is None:
            return False
        return self.min() <= other.min() and other.max() <= self.max()

    def as_constraint(self) -> Constraints:
        """
        Converts the distribution to a set of constraints.
        """
        return Constraints(
            rules=[
                (self.name, "le", self.high),
                (self.name, "ge", self.low),
                (self.name, "dtype", "float"),
            ]
        )

    def min(self) -> Any:
        """
        Returns the minimum value of the distribution.
        """
        return self.low

    def max(self) -> Any:
        """
        Returns the maximum value of the distribution.
        """
        return self.high

    def dtype(self) -> str:
        """
        Returns the data type of the distribution.
        """
        return "float"


class LogDistribution(FloatDistribution):
    low: float = np.finfo(np.float64).tiny
    high: float = np.finfo(np.float64).max

    def get(self) -> List[Any]:
        return [self.name, self.low, self.high]

    def sample(self, count: int = 1) -> Any:
        msamples = self.sample_marginal(count)
        if msamples is not None:
            return msamples
        lo, hi = np.log2(self.low), np.log2(self.high)
        return 2.0 ** self._rng.uniform(lo, hi, count)


class IntegerDistribution(Distribution):
    """
    .. inheritance-diagram:: synthcity.plugins.core.distribution.IntegerDistribution
        :parts: 1
    """

    low: Optional[int] = Field(default=None)
    high: Optional[int] = Field(default=None)
    step: int = Field(default=1)
    _is_constant: bool = PrivateAttr(False)

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @model_validator(mode="after")
    @classmethod
    def validate_and_initialize(
        cls, model: "IntegerDistribution"
    ) -> "IntegerDistribution":
        """
        Validates and initializes the distribution.
        Sets '_is_constant' based on whether 'low' equals 'high'.
        Initializes 'marginal_distribution' based on 'data' if provided.
        """
        if model.data is not None:
            # Initialize marginal_distribution based on data
            model.marginal_distribution = model.data.value_counts(
                normalize=True
            ).sort_index()
            model.low = int(model.data.min())
            model.high = int(model.data.max())
        elif model.marginal_distribution is not None:
            # Infer 'low' and 'high' from the marginal distribution's index
            model.low = int(model.marginal_distribution.index.min())
            model.high = int(model.marginal_distribution.index.max())
        else:
            # Ensure 'low' and 'high' are provided
            if model.low is None or model.high is None:
                raise ValueError(
                    "IntegerDistribution requires 'low' and 'high' values if 'data' or 'marginal_distribution' is not provided."
                )

        # Validate that low <= high
        if model.low > model.high:
            raise ValueError(
                f"Invalid range for '{model.name}': low ({model.low}) cannot be greater than high ({model.high})."
            )

        # Set _is_constant based on low == high
        model._is_constant = model.low == model.high

        # Ensure that low and high are finite integers
        if not np.isfinite(model.low) or not np.isfinite(model.high):
            raise ValueError(
                f"Invalid range for '{model.name}': low or high is not finite (low={model.low}, high={model.high})."
            )

        # Ensure that 'step' is a positive integer
        if model.step <= 0:
            raise ValueError("'step' must be a positive integer.")

        # Adjust 'low' and 'high' to be compatible with 'step'
        model.low = model.low - ((model.low - (model.low % model.step)) % model.step)
        model.high = model.high - (
            (model.high - (model.high % model.step)) % model.step
        )

        # Re-validate after adjustment
        if model.low > model.high:
            raise ValueError(
                f"After adjusting with step, invalid range for '{model.name}': low ({model.low}) cannot be greater than high ({model.high})."
            )

        return model

    def sample(self, count: int = 1) -> Any:
        """
        Samples values from the distribution.
        If the distribution is constant, returns an array filled with the constant value.
        Otherwise, samples based on the marginal distribution or uniform sampling.
        """
        if self._is_constant:
            if self.low is None:
                raise ValueError(
                    "Cannot sample: 'low' is None for a constant distribution."
                )
            samples = np.full(count, self.low)
        else:
            if self.low is None or self.high is None:
                raise ValueError("Cannot sample: 'low' or 'high' is None.")
            if (
                self.sampling_strategy == "marginal"
                and self.marginal_distribution is not None
            ):
                # Sample based on marginal distribution
                return self._rng.choice(
                    self.marginal_distribution.index,
                    size=count,
                    p=self.marginal_distribution.values,
                )
            else:
                if self.low is None or self.high is None:
                    raise ValueError(
                        "Cannot sample based on uniform distribution: low or high is not provided."
                    )
                # Proceed with uniform sampling
                possible_values = np.arange(self.low, self.high + 1, self.step)
                samples = self._rng.choice(possible_values, size=count)
        return samples

    def get(self) -> List[Any]:
        """
        Returns the metadata of the distribution.
        """
        return [self.name, self.low, self.high, self.step]

    def has(self, val: Any) -> bool:
        """
        Checks if a value is within the distribution's range.
        """
        return self.low <= val <= self.high

    def includes(self, other: "Distribution") -> bool:
        """
        Checks if another distribution is entirely within this distribution.
        """
        if self.min() is None or self.max() is None:
            return False
        if other.min() is None or other.max() is None:
            return False
        return self.min() <= other.min() and other.max() <= self.max()

    def as_constraint(self) -> Constraints:
        """
        Converts the distribution to a set of constraints.
        """
        rules: List[Rule] = []
        if self.low is not None:
            rules.append((self.name, "ge", self.low))
        if self.high is not None:
            rules.append((self.name, "le", self.high))
        rules.append((self.name, "dtype", "int"))
        return Constraints(rules=rules)

    def min(self) -> Any:
        """
        Returns the minimum value of the distribution.
        """
        return self.low

    def max(self) -> Any:
        """
        Returns the maximum value of the distribution.
        """
        return self.high

    def dtype(self) -> str:
        """
        Returns the data type of the distribution.
        """
        return "int"


class IntLogDistribution(IntegerDistribution):
    low: int = Field(default=1)
    high: int = Field(default=np.iinfo(np.int64).max)

    @field_validator("step", mode="before")
    @classmethod
    def _validate_step(cls: Any, v: int, values: ValidationInfo) -> int:
        if v != 1:
            raise ValueError("Step must be 1 for IntLogDistribution")
        return v

    def get(self) -> List[Any]:
        return [self.name, self.low, self.high]

    def sample(self, count: int = 1) -> Any:
        msamples = self.sample_marginal(count)
        if msamples is not None:
            return msamples
        lo, hi = np.log2(self.low), np.log2(self.high)
        samples = 2.0 ** self._rng.uniform(lo, hi, count)
        return samples.astype(int)


class DatetimeDistribution(Distribution):
    """
    .. inheritance-diagram:: synthcity.plugins.core.distribution.DatetimeDistribution
        :parts: 1
    """

    low: Optional[datetime] = Field(default=None)
    high: Optional[datetime] = Field(default=None)
    step: timedelta = Field(default=timedelta(microseconds=1))
    offset: timedelta = Field(default=timedelta(seconds=120))
    _is_constant: bool = PrivateAttr(False)  # Correctly named with leading underscore

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @model_validator(mode="after")
    @classmethod
    def validate_low_high(cls, model: "DatetimeDistribution") -> "DatetimeDistribution":
        """
        Validates that 'low' is less than or equal to 'high'.
        Sets '_is_constant' based on whether 'low' equals 'high'.
        """
        if model.marginal_distribution is not None:
            # Infer 'low' and 'high' from the marginal distribution's index
            model.low = model.marginal_distribution.index.min()
            model.high = model.marginal_distribution.index.max()
        else:
            # If 'marginal_distribution' is not provided, ensure 'low' and 'high' are set
            if model.low is None or model.high is None:
                if model.data is not None:
                    model.low = model.data.min()
                    model.high = model.data.max()
                else:
                    # Set default finite datetime values if not provided
                    model.low = datetime.fromtimestamp(0, timezone.utc)
                    model.high = datetime.now()
        if model.low is None or model.high is None:
            raise ValueError(
                "DatetimeDistribution requires 'low' and 'high' values if 'data' or 'marginal_distribution' is not provided."
            )
        # Validate that low <= high
        if model.low > model.high:
            raise ValueError(
                f"Invalid range for {model.name}: low ({model.low}) cannot be greater than high ({model.high})."
            )

        # Set _is_constant based on low == high
        model._is_constant = model.low == model.high

        # Ensure that low and high are valid datetime objects
        if not isinstance(model.low, datetime) or not isinstance(model.high, datetime):
            raise ValueError(
                f"Invalid range for {model.name}: low or high is not a valid datetime object (low={model.low}, high={model.high})."
            )

        # Ensure that 'step' is positive and non-zero
        if model.step.total_seconds() <= 0:
            raise ValueError("'step' must be a positive timedelta.")

        return model

    def sample(self, count: int = 1) -> Any:
        """
        Samples datetime values from the distribution.
        If the distribution is constant, returns a list filled with the constant datetime value.
        Otherwise, samples based on the specified sampling strategy.
        """
        if self._is_constant:
            if self.low is None:
                raise ValueError(
                    "Cannot sample constant datetime distribution: low is not provided."
                )
            samples = [self.low for _ in range(count)]
        else:
            if self.low is None or self.high is None:
                raise ValueError(
                    "Cannot sample datetime distribution: low or high is not provided."
                )
            if self.sampling_strategy in ["marginal", "uniform"]:
                msamples = self.sample_marginal(count)
                if msamples is not None:
                    return msamples
                if self.low is None or self.high is None:
                    raise ValueError(
                        "Cannot sample based on marginal distribution: low or high is not provided."
                    )
                total_seconds = (self.high - self.low).total_seconds()
                step_seconds = self.step.total_seconds()
                steps = int(total_seconds / step_seconds)
                step_indices = self._rng.integers(0, steps + 1, count)
                samples = [self.low + self.step * int(s) for s in step_indices]
            else:
                raise ValueError(
                    f"Unsupported sampling strategy '{self.sampling_strategy}'."
                )
        return samples

    def get(self) -> List[Any]:
        """
        Returns the metadata of the distribution.
        """
        return [self.name, self.low, self.high, self.step, self.offset]

    def has(self, val: datetime) -> bool:
        """
        Checks if a datetime value is within the distribution's range.
        """
        if self.low is None or self.high is None:
            raise ValueError("Cannot determine 'has' because 'low' or 'high' is None.")
        return self.low <= val <= self.high

    def includes(self, other: "Distribution") -> bool:
        """
        Checks if another datetime distribution is entirely within this distribution, considering the offset.
        """
        if self.low is None or self.high is None:
            return False
        if other.min() is None or other.max() is None:
            return False
        return (
            self.low - self.offset <= other.min()
            and other.max() <= self.high + self.offset
        )

    def as_constraint(self) -> Constraints:
        """
        Converts the distribution to a set of constraints.
        """
        return Constraints(
            rules=[
                (self.name, "le", self.high),
                (self.name, "ge", self.low),
                (self.name, "dtype", "datetime"),
            ]
        )

    def min(self) -> Optional[datetime]:
        """
        Returns the minimum datetime value of the distribution.
        """
        return self.low

    def max(self) -> Optional[datetime]:
        """
        Returns the maximum datetime value of the distribution.
        """
        return self.high

    def dtype(self) -> str:
        """
        Returns the data type of the distribution.
        """
        return "datetime"


class PassThroughDistribution(Distribution):
    """
    .. inheritance-diagram:: synthcity.plugins.core.distribution.PassThroughDistribution
        :parts: 1
    """

    data: pd.Series
    _dtype: str = PrivateAttr("")

    def setup_distribution(self) -> None:
        if self.data is None:
            raise ValueError("'data' must be provided for PassThroughDistribution.")

        # No additional attributes to set up since 'data' is used directly
        # Optionally, store the data type for dtype method
        self._dtype = str(self.data.dtype)

    def sample(self, count: int = 1) -> Any:
        msamples = self.sample_marginal(count)
        if msamples is not None:
            return msamples
        return self.data.sample(
            n=count, replace=True, random_state=self.random_state
        ).values

    def as_constraint(self) -> Constraints:
        # No constraints needed for pass-through columns
        return Constraints(rules=[])

    def get(self) -> List[Any]:
        # Return the unique values or any relevant info
        return [self.name]

    def has(self, val: Any) -> bool:
        # Check if the value exists in the data
        return val in self.data.values

    def includes(self, other: "Distribution") -> bool:
        # Since we are passing through values, we can define includes as checking if all values in other are in self.data
        if isinstance(other, PassThroughDistribution):
            return set(other.data.unique()).issubset(set(self.data.unique()))
        else:
            return False

    def min(self) -> Any:
        return self.data.min()

    def max(self) -> Any:
        return self.data.max()

    def dtype(self) -> str:
        return str(self.data.dtype)


def constraint_to_distribution(constraints: Constraints, feature: str) -> Distribution:
    """Infer Distribution from Constraints.

    Args:
        constraints: Constraints
            The Constraints on features.
        feature: str
            The name of the feature in question.

    Returns:
        The inferred Distribution.
    """
    dist_name, dist_args = constraints.feature_params(feature)

    if dist_name == "categorical":
        dist_template = CategoricalDistribution
    elif dist_name == "integer":
        dist_template = IntegerDistribution
    elif dist_name == "datetime":
        dist_template = DatetimeDistribution
    else:
        dist_template = FloatDistribution

    return dist_template(name=feature, **dist_args)

schema.py:

# stdlib
from typing import Any, Dict, Generator, List, Optional, Union

# third party
import pandas as pd
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    field_validator,
    model_validator,
    validate_arguments,
)

# synthcity absolute
import synthcity.logger as log
from synthcity.plugins.core.constraints import Constraints
from synthcity.plugins.core.dataloader import DataLoader, GenericDataLoader
from synthcity.plugins.core.distribution import (
    CategoricalDistribution,
    DatetimeDistribution,
    Distribution,
    FloatDistribution,
    IntegerDistribution,
    PassThroughDistribution,
)


class Schema(BaseModel):
    """
    Utility class for defining the schema of a Dataset.

    Constructor Args:
        domain: Dict
            A dictionary of feature_name: Distribution.
        sampling_strategy: str
            Taking value of "marginal" (default) or "uniform" (for debugging).
        protected_cols: List[str]
            List of columns that are exempt from distributional constraints (e.g. ID column)
        random_state: int
            Random seed (default 0)
        data: Any
            (Optional) the data set
    """

    sampling_strategy: str = Field(default="marginal")
    protected_cols: List[str] = []
    random_state: int = Field(default=0)
    domain: Dict = Field(default_factory=dict)

    data: Optional[Union[DataLoader, pd.DataFrame]] = Field(default=None, exclude=True)

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @field_validator("data", mode="before")
    @classmethod
    def validate_data(cls, v: Any) -> Optional[DataLoader]:
        if v is not None:
            if isinstance(v, pd.DataFrame):
                return GenericDataLoader(v)
            elif isinstance(v, DataLoader):
                return v
            else:
                raise ValueError(
                    f"Invalid data type for 'data': {type(v)}. Expected DataLoader or pandas DataFrame."
                )
        return v

    @model_validator(mode="after")
    @classmethod 
    def initialize_domain(cls, model: "Schema") -> "Schema":
        if model.data is not None:
            X = model.data.dataframe()
            model.domain = model._infer_domain(
                X,
                sampling_strategy=model.sampling_strategy,
                random_state=model.random_state,
            )
            # Remove 'data' attribute from the model
            del model.__dict__["data"]
            if "data" in model.__fields_set__:
                model.__fields_set__.remove("data")
        return model

    @validate_arguments
    def get(self, feature: str) -> Distribution:
        """Get the Distribution of a feature.

        Args:
            feature: str. the feature name

        Returns:
            The feature distribution
        """
        if feature not in self.domain:
            raise ValueError(f"invalid feature {feature}")

        return self.domain[feature]

    @validate_arguments
    def __getitem__(self, key: str) -> Distribution:
        """Get the Distribution of a feature.

        Args:
            feature: str. the feature name

        Returns:
            The feature distribution
        """
        return self.get(key)

    def __iter__(self) -> Generator:
        """Iterate the features distribution"""
        for x in self.domain:
            yield x

    def __len__(self) -> int:
        """Get the number of features"""
        return len(self.domain)

    def includes(self, other: "Schema") -> bool:
        """Test if another schema is included in the local one."""
        for feature in other:
            if feature in self.protected_cols:
                continue
            if feature not in self.domain:
                return False

            if not self[feature].includes(other[feature]):
                return False

        return True

    def features(self) -> List:
        return list(self.domain.keys())

    def sample(self, count: int) -> pd.DataFrame:
        data = {}
        for col, dist in self.domain.items():
            samples = dist.sample(count)
            data[col] = samples
        return pd.DataFrame(data)

    def adapt_dtypes(self, X: pd.DataFrame) -> pd.DataFrame:
        """Applying the data type to a new data frame

        Args:
            X: pd.DataFrame
                A new data frame to be adapted.

        Returns:
            A data frame whose data types are coerced to be the same with the Schema.
            If the data frame contains new features, these will be retained as is.
        """
        for feature in self.domain:
            if feature not in X.columns:
                continue
            X[feature] = X[feature].astype(
                self.domain[feature].dtype(), errors="ignore"
            )

        return X

    def as_constraints(self) -> Constraints:
        rules = []
        for feature, dist in self.domain.items():
            rules.extend(dist.as_constraint().rules)
        return Constraints(rules=rules)

    @classmethod
    def from_constraints(cls, constraints: Constraints) -> "Schema":
        domain: Dict = {}
        feature_params: Dict = {}

        # Collect constraint information
        for feature, op, value in constraints.rules:
            if feature not in feature_params:
                feature_params[feature] = {
                    "name": feature,
                    "random_state": None,
                    "low": None,
                    "high": None,
                    "dtype": "float",  # Default to 'float' if not specified
                    "choices": [],
                }

            params = feature_params[feature]

            if op in ["ge", ">="]:
                if params["low"] is None or value > params["low"]:
                    params["low"] = value
            elif op in ["le", "<="]:
                if params["high"] is None or value < params["high"]:
                    params["high"] = value
            elif op in ["eq", "=="]:
                # For '==', set both 'low' and 'high' to value
                params["low"] = value
                params["high"] = value
            elif op in ["in", "isin"]:
                if isinstance(value, list):
                    params["choices"].extend(value)
                else:
                    params["choices"].append(value)
            elif op == "dtype":
                params["dtype"] = value
            else:
                # Handle other operators if necessary
                pass

        # Create distribution objects
        for feature, params in feature_params.items():
            dtype = params["dtype"]
            if dtype == "float":
                if params["low"] is None or params["high"] is None:
                    raise ValueError(
                        f"Cannot create FloatDistribution for '{feature}' without 'low' and 'high' values."
                    )
                domain[feature] = FloatDistribution(
                    name=params["name"],
                    random_state=params["random_state"],
                    low=params["low"],
                    high=params["high"],
                )
            elif dtype == "int":
                if params["low"] is None or params["high"] is None:
                    raise ValueError(
                        f"Cannot create IntegerDistribution for '{feature}' without 'low' and 'high' values."
                    )
                domain[feature] = IntegerDistribution(
                    name=params["name"],
                    random_state=params["random_state"],
                    low=int(params["low"]),
                    high=int(params["high"]),
                    step=1,  # Default step to 1 or adjust as needed
                )
            elif dtype in ["category", "object"]:
                choices = params.get("choices")
                if choices is None or not choices:
                    raise ValueError(
                        f"Cannot create CategoricalDistribution for '{feature}' without 'choices'."
                    )
                domain[feature] = CategoricalDistribution(
                    name=params["name"],
                    random_state=params["random_state"],
                    choices=list(set(choices)),
                )
            else:
                raise ValueError(
                    f"Unsupported dtype '{dtype}' for feature '{feature}'."
                )

        return cls(domain=domain)

    def _infer_domain(
        self,
        X: pd.DataFrame,
        sampling_strategy: str,
        random_state: int,
    ) -> Dict[str, Distribution]:
        feature_domain: Dict[str, Distribution] = {}

        for idx, col in enumerate(X.columns):
            col_random_state = random_state + idx + 1  # Ensure unique seeds

            try:
                if sampling_strategy == "marginal":
                    if col in self.protected_cols:
                        feature_domain[col] = PassThroughDistribution(
                            name=col,
                            data=X[col],
                            random_state=col_random_state,
                        )
                        continue

                    is_categorical = pd.api.types.is_categorical_dtype(X[col])
                    is_object = X[col].dtype == object
                    is_bool = pd.api.types.is_bool_dtype(X[col])
                    is_integer = pd.api.types.is_integer_dtype(X[col])
                    is_float = pd.api.types.is_float_dtype(X[col])
                    is_datetime = pd.api.types.is_datetime64_any_dtype(X[col])

                    if is_categorical or is_object or is_bool:
                        feature_domain[col] = CategoricalDistribution(
                            name=col,
                            data=X[col],
                            random_state=col_random_state,
                        )
                    elif is_integer:
                        feature_domain[col] = IntegerDistribution(
                            name=col,
                            data=X[col],
                            random_state=col_random_state,
                        )
                    elif is_float:
                        feature_domain[col] = FloatDistribution(
                            name=col,
                            data=X[col],
                            random_state=col_random_state,
                        )
                    elif is_datetime:
                        feature_domain[col] = DatetimeDistribution(
                            name=col,
                            data=X[col],
                            random_state=col_random_state,
                        )
                    else:
                        raise ValueError(
                            f"Unsupported data type for column '{col}' with dtype {X[col].dtype}"
                        )
                elif sampling_strategy == "uniform":

                    is_categorical = pd.api.types.is_categorical_dtype(X[col])
                    is_object = X[col].dtype == object
                    is_bool = pd.api.types.is_bool_dtype(X[col])
                    is_integer = pd.api.types.is_integer_dtype(X[col])
                    is_float = pd.api.types.is_float_dtype(X[col])
                    is_datetime = pd.api.types.is_datetime64_any_dtype(X[col])

                    if (
                        pd.api.types.is_categorical_dtype(X[col])
                        or X[col].dtype == object
                        or pd.api.types.is_bool_dtype(X[col])
                    ):
                        feature_domain[col] = CategoricalDistribution(
                            name=col,
                            choices=list(X[col].unique()),
                            random_state=col_random_state,
                            sampling_strategy=sampling_strategy,
                        )
                    elif pd.api.types.is_integer_dtype(X[col]):
                        feature_domain[col] = IntegerDistribution(
                            name=col,
                            low=X[col].min(),
                            high=X[col].max(),
                            random_state=col_random_state,
                            sampling_strategy=sampling_strategy,
                        )
                    elif pd.api.types.is_float_dtype(X[col]):
                        feature_domain[col] = FloatDistribution(
                            name=col,
                            low=X[col].min(),
                            high=X[col].max(),
                            random_state=col_random_state,
                            sampling_strategy=sampling_strategy,
                        )
                    elif pd.api.types.is_datetime64_any_dtype(X[col]):
                        feature_domain[col] = DatetimeDistribution(
                            name=col,
                            low=X[col].min(),
                            high=X[col].max(),
                            random_state=col_random_state,
                            sampling_strategy=sampling_strategy,
                        )
                else:
                    raise ValueError(
                        f"Unsupported sampling strategy '{sampling_strategy}'"
                    )
            except Exception as e:
                log.error(f"Exception occurred while processing column '{col}': {e}")
                raise
        return feature_domain

[woznicak]

@M-Zeeb, I dealt with the same problem as @atulkakrana . Thanks for sharing this solution — it looks great!

I tried running it, but I’m running into some issues that might be related to library versions.
Would you mind sharing the list of packages (and their versions) from your environment?

You can get that by running:

pip freeze > requirements.txt

or simply:

pip list

That would really help me reproduce your setup and troubleshoot the problem. Thanks a lot!

[M-Zeeb]

Just for consideration, I'm running it within SynthCraft.
I also adapted the rms_norm.py under miniforge3/envs/SynthCraft/lib/python3.9/site-packages/opacus/grad_sample/ to fix the torch.nn error with PyTorch 2.2.2. However, full disclosure, I just used AI here for a quickfix and did not test it.

import torch
import torch.nn as nn
import torch.nn.functional as F
from opacus.utils.tensor_utils import sum_over_all_but_batch_and_last_n

from .utils import register_grad_sampler

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-8):
        """
        RMSNorm implementation compatible with PyTorch 2.2.2
        Args:
            dim (int): number of features in the last dimension
            eps (float): small value to avoid division by zero
        """
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        """
        x: tensor of shape (..., dim)
        returns: normalized tensor of same shape as x
        """
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
        return x / rms * self.scale

@register_grad_sampler(RMSNorm)
def compute_rms_norm_grad_sample(
    layer: RMSNorm,
    activations: List[torch.Tensor],
    backprops: torch.Tensor,
) -> Dict[nn.Parameter, torch.Tensor]:
    """
    Computes per sample gradients for RMSNorm

    Args:
        layer: Layer
        activations: Activations
        backprops: Backpropagations
    """
    activations = activations[0]
    ret = {}
    if layer.weight.requires_grad:
        ret[layer.weight] = sum_over_all_but_batch_and_last_n(
            F.rms_norm(activations, layer.normalized_shape, eps=layer.eps) * backprops,
            layer.weight.dim(),
        )
    return ret

Package                      Version                              Editable project location
---------------------------- ------------------------------------ ----------------------------------------------------
absl-py                      2.3.1
accelerate                   1.10.1
aiohappyeyeballs             2.6.1
aiohttp                      3.13.0
aiosignal                    1.4.0
alembic                      1.16.5
altair                       5.5.0
annotated-types              0.7.0
anyio                        4.11.0
arfpy                        0.1.1
asttokens                    3.0.0
async-timeout                5.0.1
attrs                        25.4.0
autograd                     1.8.0
autograd-gamma               0.5.0
be_great                     0.0.9
blinker                      1.9.0
blis                         1.2.1
Brotli                       1.1.0
cachetools                   6.2.1
catalogue                    2.0.10
certifi                      2025.10.5
cffi                         2.0.0
charset-normalizer           3.4.4
click                        8.1.8
climb-ai                     0.0.post1.dev95+gc01f0e247.d20251016
cloudpathlib                 0.23.0
cloudpickle                  3.1.1
colorlog                     6.9.0
confection                   0.1.5
contourpy                    1.3.0
cryptography                 46.0.3
cssselect2                   0.8.0
cycler                       0.12.1
cymem                        2.0.11
datasets                     4.2.0
decaf-synthetic-data         0.1.7
decorator                    5.2.1
dgl                          2.0.0
dill                         0.4.0
distro                       1.9.0
exceptiongroup               1.3.0
executing                    2.2.0
fastai                       2.7.19
fastcore                     1.7.29
fastdownload                 0.0.7
fastprogress                 1.0.3
feather-format               0.4.1
fflows                       0.0.3
filelock                     3.19.1
fonttools                    4.59.1
formulaic                    1.2.1
frozenlist                   1.8.0
fsspec                       2025.9.0
geomloss                     0.2.6
gitdb                        4.0.12
GitPython                    3.1.45
google-ai-generativelanguage 0.6.15
google-api-core              2.26.0
google-api-python-client     2.184.0
google-auth                  2.41.1
google-auth-httplib2         0.2.0
google-generativeai          0.8.5
googleapis-common-protos     1.70.0
GPUtil                       1.4.0
greenlet                     3.2.4
grpcio                       1.75.1
grpcio-status                1.71.2
h11                          0.16.0
h5py                         3.14.0
hf-xet                       1.1.10
httpcore                     1.0.9
httplib2                     0.31.0
httpx                        0.28.1
huggingface-hub              0.35.3
idna                         3.11
imbalanced-learn             0.12.4
importlib_metadata           8.7.0
importlib_resources          6.5.2
inflate64                    1.0.3
interface-meta               1.3.0
ipython                      8.18.1
jedi                         0.19.2
Jinja2                       3.1.6
jiter                        0.11.0
joblib                       1.5.2
jsonschema                   4.25.1
jsonschema-specifications    2025.9.1
keopscore                    2.3
kiwisolver                   1.4.7
langcodes                    3.5.0
language_data                1.3.0
lifelines                    0.29.0
lightning-utilities          0.15.2
llvmlite                     0.43.0
loguru                       0.7.3
Mako                         1.3.10
marisa-trie                  1.3.1
Markdown                     3.9
markdown-it-py               3.0.0
MarkupSafe                   3.0.3
matplotlib                   3.9.4
matplotlib-inline            0.1.7
mdurl                        0.1.2
monai                        1.4.0
mpmath                       1.3.0
multidict                    6.7.0
multiprocess                 0.70.16
multivolumefile              0.2.3
munkres                      1.1.4
murmurhash                   1.0.13
narwhals                     2.8.0
networkx                     2.8.8
nflows                       0.14
numba                        0.60.0
numpy                        1.26.4
nutree                       1.1.0
nvidia-cublas-cu12           12.1.3.1
nvidia-cuda-cupti-cu12       12.1.105
nvidia-cuda-nvrtc-cu12       12.1.105
nvidia-cuda-runtime-cu12     12.1.105
nvidia-cudnn-cu12            8.9.2.26
nvidia-cufft-cu12            11.0.2.54
nvidia-curand-cu12           10.3.2.106
nvidia-cusolver-cu12         11.4.5.107
nvidia-cusparse-cu12         12.1.0.106
nvidia-nccl-cu12             2.19.3
nvidia-nvjitlink-cu12        12.9.86
nvidia-nvtx-cu12             12.1.105
opacus                       1.5.4
openai                       2.3.0
opt_einsum                   3.4.0
optuna                       4.5.0
packaging                    25.0
pandas                       2.3.3
parso                        0.8.4
patsy                        1.0.1
pdfminer.six                 20250506
pdfplumber                   0.11.7
pexpect                      4.9.0
pgmpy                        0.1.26
pickleshare                  0.7.5
pillow                       11.3.0
pip                          25.2
plotly                       6.3.1
preshed                      3.0.10
prompt_toolkit               3.0.51
propcache                    0.4.1
proto-plus                   1.26.1
protobuf                     5.29.5
psutil                       7.1.0
ptyprocess                   0.7.0
pure_eval                    0.2.3
py7zr                        1.0.0
pyarrow                      21.0.0
pyasn1                       0.6.1
pyasn1_modules               0.4.2
pybcj                        1.0.6
pybind11                     3.0.1
pycox                        0.3.0
pycparser                    2.22
pycryptodomex                3.23.0
pydantic                     2.12.2
pydantic_core                2.41.4
pydeck                       0.9.1
pydyf                        0.11.0
Pygments                     2.19.2
pykeops                      2.3
pyparsing                    3.2.5
pypdfium2                    4.30.0
pyphen                       0.17.2
pyppmd                       1.2.0
python-dateutil              2.9.0.post0
python-dotenv                1.1.1
pytorch-lightning            2.5.5
pyts                         0.13.0
pytz                         2025.2
PyYAML                       6.0.3
pyzstd                       0.18.0
redis                        6.4.0
referencing                  0.36.2
regex                        2025.9.18
requests                     2.32.5
rich                         14.2.0
rpds-py                      0.27.1
rsa                          4.9.1
ruamel.yaml                  0.18.15
ruamel.yaml.clib             0.2.14
safetensors                  0.6.2
SciencePlots                 2.1.1
scikit-learn                 1.6.1
scipy                        1.13.1
seaborn                      0.13.2
setuptools                   80.9.0
shap                         0.49.1
shellingham                  1.5.4
six                          1.17.0
slicer                       0.0.8
smart_open                   7.3.1
smmap                        5.0.2
sniffio                      1.3.1
spacy                        3.8.7
spacy-legacy                 3.0.12
spacy-loggers                1.0.5
SQLAlchemy                   2.0.44
srsly                        2.5.1
stack_data                   0.6.3
statsmodels                  0.14.5
streamlit                    1.50.0
streamlit-antd-components    0.3.2
streamlit-js                 1.0.8
sympy                        1.14.0
synthcity                    0.2.12
tenacity                     9.1.2
tensorboard                  2.20.0
tensorboard-data-server      0.7.2
texttable                    1.7.0
thinc                        8.3.4
threadpoolctl                3.6.0
tiktoken                     0.12.0
tinycss2                     1.4.0
tinydb                       4.8.2
tinydb-serialization         2.2.0
tinyhtml5                    2.0.0
tokenizers                   0.22.1
toml                         0.10.2
tomli                        2.3.0
torch                        2.2.2
torch-geometric              2.6.1
torch_scatter                2.1.2
torch_sparse                 0.6.18
torchdata                    0.11.0
torchmetrics                 1.8.2
torchtext                    0.17.2
torchtuples                  0.2.2
torchvision                  0.17.2
tornado                      6.5.2
tqdm                         4.67.1
traitlets                    5.14.3
transformers                 4.57.1
triton                       2.2.0
tsai                         0.4.0
typer                        0.19.2
typing_extensions            4.14.1
typing-inspection            0.4.2
tzdata                       2025.2
unicodedata2                 16.0.0
uritemplate                  4.2.0
urllib3                      2.5.0
wasabi                       1.1.3
watchdog                     6.0.0
wcwidth                      0.2.13
weasel                       0.4.1
weasyprint                   65.1
webencodings                 0.5.1
Werkzeug                     3.1.3
wheel                        0.45.1
wrapt                        1.17.3
xgboost                      2.1.4
xgbse                        0.3.3
xxhash                       3.6.0
yarl                         1.22.0
zipp                         3.23.0
zopfli                       0.2.3.post1