kumoai 2.14.0.dev202512181731__cp312-cp312-macosx_11_0_arm64.whl → 2.14.0.dev202601041732__cp312-cp312-macosx_11_0_arm64.whl

kumoai/experimental/rfm/task_table.py

@@ -0,0 +1,247 @@
+import copy
+from collections.abc import Sequence
+
+import pandas as pd
+from kumoapi.task import TaskType
+from kumoapi.typing import Dtype, Stype
+from typing_extensions import Self
+
+from kumoai.experimental.rfm.base import Column
+from kumoai.experimental.rfm.infer import contains_timestamp, infer_dtype
+
+
+class TaskTable:
+    r"""A :class:`TaskTable` fully specifies the task, *i.e.* its context and
+    prediction examples with entity IDs, targets and timestamps.
+
+    Args:
+        task_type: The task type.
+        context_df: The data frame holding context examples.
+        pred_df: The data frame holding prediction examples.
+        entity_table_name: The entity table to predict for. For link prediction
+            tasks, needs to hold both entity and target table names.
+        entity_column: The name of the entity column.
+        target_column: The name of the target column.
+        time_column: The name of the time column, if it exists.
+    """
+    def __init__(
+        self,
+        task_type: TaskType,
+        context_df: pd.DataFrame,
+        pred_df: pd.DataFrame,
+        entity_table_name: str | Sequence[str],
+        entity_column: str,
+        target_column: str,
+        time_column: str | None = None,
+    ) -> None:
+
+        task_type = TaskType(task_type)
+        if task_type not in {  # Currently supported task types:
+                TaskType.BINARY_CLASSIFICATION,
+                TaskType.MULTICLASS_CLASSIFICATION,
+                TaskType.REGRESSION,
+                TaskType.TEMPORAL_LINK_PREDICTION,
+        }:
+            raise ValueError  # TODO
+        self._task_type = task_type
+
+        # TODO Check dfs (unify from local table)
+        self._context_df = context_df.copy(deep=False)
+        self._pred_df = pred_df.copy(deep=False)
+
+        self._dtype_dict: dict[str, Dtype] = {
+            column_name: infer_dtype(context_df[column_name])
+            for column_name in context_df.columns
+        }
+
+        self._entity_table_names: tuple[str] | tuple[str, str]
+        if isinstance(entity_table_name, str):
+            self._entity_table_names = (entity_table_name, )
+        elif len(entity_table_name) == 1:
+            self._entity_table_names = (entity_table_name[0], )
+        elif len(entity_table_name) == 2:
+            self._entity_table_names = (
+                entity_table_name[0],
+                entity_table_name[1],
+            )
+        else:
+            raise ValueError  # TODO
+
+        self._entity_column: str = ''
+        self._target_column: str = ''
+        self._time_column: str | None = None
+
+        self.entity_column = entity_column
+        self.target_column = target_column
+        if time_column is not None:
+            self.time_column = time_column
+
+    @property
+    def task_type(self) -> TaskType:
+        r"""The task type."""
+        return self._task_type
+
+    def narrow(self, start: int, length: int) -> Self:
+        r"""Returns a new :class:`TaskTable` that holds a narrowed version of
+        prediction examples.
+
+        Args:
+            start: Index of the prediction examples to start narrowing.
+            length: Length of the prediction examples.
+        """
+        out = copy.copy(self)
+        df = out._pred_df.iloc[start:start + length].reset_index(drop=True)
+        out._pred_df = df
+        return out
+
+    # Entity column ###########################################################
+
+    @property
+    def entity_table_name(self) -> str:
+        return self._entity_table_names[0]
+
+    @property
+    def entity_table_names(self) -> tuple[str] | tuple[str, str]:
+        return self._entity_table_names
+
+    @property
+    def entity_column(self) -> Column:
+        return Column(
+            name=self._entity_column,
+            expr=None,
+            dtype=self._dtype_dict[self._entity_column],
+            stype=Stype.ID,
+        )
+
+    @entity_column.setter
+    def entity_column(self, name: str) -> None:
+        if name in {self._target_column, self._time_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if name not in self._pred_df:
+            raise ValueError  # TODO
+        if not Stype.ID.supports_dtype(self._dtype_dict[name]):
+            raise ValueError  # TODO
+
+        self._entity_column = name
+
+    # Target column ###########################################################
+
+    @property
+    def _target_stype(self) -> Stype:
+        if self.task_type in {
+                TaskType.BINARY_CLASSIFICATION,
+                TaskType.MULTICLASS_CLASSIFICATION,
+        }:
+            return Stype.categorical
+        if self.task_type in {TaskType.REGRESSION}:
+            return Stype.numerical
+        if self.task_type.is_link_pred:
+            return Stype.multicategorical
+        raise ValueError
+
+    @property
+    def target_column(self) -> Column:
+        return Column(
+            name=self._target_column,
+            expr=None,
+            dtype=self._dtype_dict[self._target_column],
+            stype=self._target_stype,
+        )
+
+    @target_column.setter
+    def target_column(self, name: str) -> None:
+        if name in {self._entity_column, self._time_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if not self._target_stype.supports_dtype(self._dtype_dict[name]):
+            raise ValueError  # TODO
+
+        self._target_column = name
+
+    # Time column #############################################################
+
+    def has_time_column(self) -> bool:
+        r"""Returns ``True`` if this task has a time column; ``False``
+        otherwise.
+        """
+        return self._time_column is not None
+
+    @property
+    def time_column(self) -> Column | None:
+        r"""The time column of this task.
+
+        The getter returns the time column of this task, or ``None`` if no
+        such time column is present.
+
+        The setter sets a column as a time column for this task, and raises a
+        :class:`ValueError` if the time column has a non-timestamp compatible
+        data type or if the column name does not match a column in the data
+        frame.
+        """
+        if self._time_column is None:
+            return None
+        return Column(
+            name=self._time_column,
+            expr=None,
+            dtype=self._dtype_dict[self._time_column],
+            stype=Stype.timestamp,
+        )
+
+    @time_column.setter
+    def time_column(self, name: str | None) -> None:
+        if name is None:
+            self._time_column = None
+            return
+
+        if name in {self._entity_column, self._target_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if name not in self._pred_df:
+            raise ValueError  # TODO
+        if not contains_timestamp(
+                ser=self._context_df[name],
+                column_name=name,
+                dtype=self._dtype_dict[name],
+        ):
+            raise ValueError  # TODO
+
+        self._time_column = name
+
+    # Metadata ################################################################
+
+    @property
+    def metadata(self) -> pd.DataFrame:
+        raise NotImplementedError
+
+    def print_metadata(self) -> None:
+        raise NotImplementedError
+
+    # Python builtins #########################################################
+
+    def __hash__(self) -> int:
+        return hash((
+            self.task_type,
+            self.entity_table_names,
+            self._entity_column,
+            self._target_column,
+            self._time_column,
+        ))
+
+    def __repr__(self) -> str:
+        if self.task_type.is_link_pred:
+            entity_table_repr = f'entity_table_names={self.entity_table_names}'
+        else:
+            entity_table_repr = f'entity_table_name={self.entity_table_name}'
+        return (f'{self.__class__.__name__}(\n'
+                f'  task_type={self.task_type},\n'
+                f'  num_context_examples={len(self._context_df)},\n'
+                f'  num_prediction_examples={len(self._pred_df)},\n'
+                f'  {entity_table_repr},\n'
+                f'  entity_column={self._entity_column},\n'
+                f'  target_column={self._target_column},\n'
+                f'  time_column={self._time_column},\n'
+                f')')

kumoai/trainer/distilled_trainer.py

@@ -0,0 +1,175 @@
+import logging
+from typing import Literal, Mapping, Optional, Union, overload
+
+from kumoapi.distilled_model_plan import DistilledModelPlan
+from kumoapi.jobs import DistillationJobRequest, DistillationJobResource
+
+from kumoai import global_state
+from kumoai.client.jobs import TrainingJobID
+from kumoai.graph import Graph
+from kumoai.pquery.training_table import TrainingTable, TrainingTableJob
+from kumoai.trainer.job import TrainingJob, TrainingJobResult
+
+logger = logging.getLogger(__name__)
+
+
+class DistillationTrainer:
+    r"""A trainer supports creating a Kumo machine learning model
+    for use in an online serving endpoint. The distllation process involes
+    training a shallow model on a :class:`~kumoai.pquery.PredictiveQuery` using
+    the embeddings generated by a base model :args:`base_training_job_id`.
+
+    Args:
+        model_plan: The distilled model plan to use for the distillation process.
+        base_training_job_id: The ID of the base training job to use for the distillation process.
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        model_plan: DistilledModelPlan,
+        base_training_job_id: TrainingJobID,
+    ) -> None:
+        self.model_plan: DistilledModelPlan = model_plan
+        self.base_training_job_id: TrainingJobID = base_training_job_id
+
+        # Cached from backend:
+        self._training_job_id: Optional[TrainingJobID] = None
+
+    # Metadata ################################################################
+
+    @property
+    def is_trained(self) -> bool:
+        r"""Returns ``True`` if this trainer instance has successfully been
+        trained (and is therefore ready for prediction); ``False`` otherwise.
+        """
+        raise NotImplementedError(
+            "Checking if a distilled trainer is trained is not "
+            "implemented yet.")
+
+    @overload
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+    ) -> TrainingJobResult:
+        pass
+
+    @overload
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: Literal[False],
+    ) -> TrainingJobResult:
+        pass
+
+    @overload
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: Literal[True],
+    ) -> TrainingJob:
+        pass
+
+    @overload
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: bool,
+    ) -> Union[TrainingJob, TrainingJobResult]:
+        pass
+
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: bool = False,
+        custom_tags: Mapping[str, str] = {},
+    ) -> Union[TrainingJob, TrainingJobResult]:
+        r"""Fits a model to the specified graph and training table, with the
+        strategy defined by :class:`DistilledTrainer`'s :obj:`model_plan`.
+
+        Args:
+            graph: The :class:`~kumoai.graph.Graph` object that represents the
+                tables and relationships that Kumo will learn from.
+            train_table: The :class:`~kumoai.pquery.TrainingTable`, or
+                in-progress :class:`~kumoai.pquery.TrainingTableJob`, that
+                represents the training data produced by a
+                :class:`~kumoai.pquery.PredictiveQuery` on :obj:`graph`.
+            non_blocking: Whether this operation should return immediately
+                after launching the training job, or await completion of the
+                training job.
+            custom_tags: Additional, customer defined k-v tags to be associated
+                with the job to be launched. Job tags are useful for grouping
+                and searching jobs.
+
+        Returns:
+            Union[TrainingJobResult, TrainingJob]:
+                If ``non_blocking=False``, returns a training job object. If
+                ``non_blocking=True``, returns a training job future object.
+        """
+        # TODO(manan, siyang): remove soon:
+        job_id = train_table.job_id
+        assert job_id is not None
+
+        train_table_job_api = global_state.client.generate_train_table_job_api
+        pq_id = train_table_job_api.get(job_id).config.pquery_id
+        assert pq_id is not None
+
+        custom_table = None
+        if isinstance(train_table, TrainingTable):
+            custom_table = train_table._custom_train_table
+
+        # NOTE the backend implementation currently handles sequentialization
+        # between a training table future and a training job; that is, if the
+        # training table future is still executing, the backend will wait on
+        # the job ID completion before executing a training job. This preserves
+        # semantics for both futures, ensures that Kumo works as expected if
+        # used only via REST API, and allows us to avoid chaining calllbacks
+        # in an ugly way here:
+        api = global_state.client.distillation_job_api
+        self._training_job_id = api.create(
+            DistillationJobRequest(
+                dict(custom_tags),
+                pquery_id=pq_id,
+                base_training_job_id=self.base_training_job_id,
+                distilled_model_plan=self.model_plan,
+                graph_snapshot_id=graph.snapshot(non_blocking=non_blocking),
+                train_table_job_id=job_id,
+                custom_train_table=custom_table,
+            ))
+
+        out = TrainingJob(job_id=self._training_job_id)
+        if non_blocking:
+            return out
+        return out.attach()
+
+    @classmethod
+    def _load_from_job(
+        cls,
+        job: DistillationJobResource,
+    ) -> 'DistillationTrainer':
+        trainer = cls(job.config.distilled_model_plan,
+                      job.config.base_training_job_id)
+        trainer._training_job_id = job.job_id
+        return trainer
+
+    @classmethod
+    def load(cls, job_id: TrainingJobID) -> 'DistillationTrainer':
+        r"""Creates a :class:`~kumoai.trainer.Trainer` instance from a training
+        job ID.
+        """
+        raise NotImplementedError(
+            "Loading a distilled trainer from a job ID is not implemented yet."
+        )
+
+    @classmethod
+    def load_from_tags(cls, tags: Mapping[str, str]) -> 'DistillationTrainer':
+        raise NotImplementedError(
+            "Loading a distilled trainer from tags is not implemented yet.")

kumoai/utils/display.py

@@ -0,0 +1,51 @@
+from collections.abc import Sequence
+
+import pandas as pd
+
+from kumoai import in_notebook, in_snowflake_notebook
+
+
+def message(msg: str) -> None:
+    msg = msg.replace("`", "'") if not in_notebook() else msg
+
+    if in_snowflake_notebook():
+        import streamlit as st
+        st.markdown(msg)
+    elif in_notebook():
+        from IPython.display import Markdown, display
+        display(Markdown(msg))
+    else:
+        print(msg)
+
+
+def title(msg: str) -> None:
+    message(f"### {msg}" if in_notebook() else f"{msg}:")
+
+
+def italic(msg: str) -> None:
+    message(f"*{msg}*" if in_notebook() else msg)
+
+
+def unordered_list(items: Sequence[str]) -> None:
+    if in_notebook():
+        msg = '\n'.join([f"- {item}" for item in items])
+    else:
+        msg = '\n'.join([f"• {item.replace('`', '')}" for item in items])
+    message(msg)
+
+
+def dataframe(df: pd.DataFrame) -> None:
+    if in_snowflake_notebook():
+        import streamlit as st
+        st.dataframe(df, hide_index=True)
+    elif in_notebook():
+        from IPython.display import display
+        try:
+            if hasattr(df.style, 'hide'):
+                display(df.style.hide(axis='index'))  # pandas=2
+            else:
+                display(df.style.hide_index())  # pandas<1.3
+        except ImportError:
+            print(df.to_string(index=False))  # missing jinja2
+    else:
+        print(df.to_string(index=False))

{kumoai-2.14.0.dev202512181731.dist-info → kumoai-2.14.0.dev202601041732.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kumoai
-Version: 2.14.0.dev202512181731
+Version: 2.14.0.dev202601041732
 Summary: AI on the Modern Data Stack
 Author-email: "Kumo.AI" <hello@kumo.ai>
 License-Expression: MIT

{kumoai-2.14.0.dev202512181731.dist-info → kumoai-2.14.0.dev202601041732.dist-info}/RECORD

@@ -1,8 +1,8 @@
 kumoai/_logging.py,sha256=U2_5ROdyk92P4xO4H2WJV8EC7dr6YxmmnM-b7QX9M7I,886
 kumoai/mixin.py,sha256=MP413xzuCqWhxAPUHmloLA3j4ZyF1tEtfi516b_hOXQ,812
-kumoai/_version.py,sha256=x_VP--ZQYmARddiqEJCsKj7I1LaFz23DUBdstrPepwI,39
+kumoai/_version.py,sha256=ZgEOgzhMS-JHEcx_StKbib3QyH7C44Vgs2zNQ7IM43A,39
 kumoai/kumolib.cpython-312-darwin.so,sha256=xQvdWHx9xmQ11y3F3ywxJv6A0sDk6D3-2fQbxSdM1z4,232576
-kumoai/__init__.py,sha256=Nn9YH_x9kAeEFn8RWbP95slZow0qFnakPZZ1WADe1hY,10843
+kumoai/__init__.py,sha256=x6Emn6VesHQz0wR7ZnbddPRYO9A5-0JTHDkzJ3Ocq6w,10907
 kumoai/formatting.py,sha256=jA_rLDCGKZI8WWCha-vtuLenVKTZvli99Tqpurz1H84,953
 kumoai/futures.py,sha256=oJFIfdCM_3nWIqQteBKYMY4fPhoYlYWE_JA2o6tx-ng,3737
 kumoai/jobs.py,sha256=NrdLEFNo7oeCYSy-kj2nAvCFrz9BZ_xrhkqHFHk5ksY,2496
@@ -11,41 +11,43 @@ kumoai/databricks.py,sha256=e6E4lOFvZHXFwh4CO1kXU1zzDU3AapLQYMxjiHPC-HQ,476
 kumoai/spcs.py,sha256=N31d7rLa-bgYh8e2J4YzX1ScxGLqiVXrqJnCl1y4Mts,4139
 kumoai/_singleton.py,sha256=UTwrbDkoZSGB8ZelorvprPDDv9uZkUi1q_SrmsyngpQ,836
 kumoai/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kumoai/experimental/rfm/graph.py,sha256=4gtGKHgLVRMNQrbXieSWtImggX5KDRf4XrLk0LcF-AA,44446
-kumoai/experimental/rfm/__init__.py,sha256=9aelcHodt2Oriw76vdEmtWrmAQ0CXTdFPrKgwVB9eKc,7124
+kumoai/experimental/rfm/relbench.py,sha256=cVsxxV3TIL3PLEoYb-8tAVW3GSef6NQAd3rxdHJL63I,2276
+kumoai/experimental/rfm/graph.py,sha256=H9lIQLDkL5zJMwEHh7PgruvMUxWsjpynXUT7gnmTTUM,46351
+kumoai/experimental/rfm/__init__.py,sha256=bW2XyYtkbdiu_iICYFF2Fu1Fx5fyGbqne6m_6c1P-fY,7016
 kumoai/experimental/rfm/sagemaker.py,sha256=6fyXO1Jd_scq-DH7kcv6JcV8QPyTbh4ceqwQDPADlZ0,4963
-kumoai/experimental/rfm/rfm.py,sha256=Yssmo-PaCfjT9hll0BKl8fahsuKpG-gViwFdKH1F3os,50247
+kumoai/experimental/rfm/rfm.py,sha256=M5Q_vasSVrAX-5ucFPrdWoTZdpE4VYxKjKZz7BccITc,49672
 kumoai/experimental/rfm/authenticate.py,sha256=G2RkRWznMVQUzvhvbKhn0bMCY7VmoNYxluz3THRqSdE,18851
+kumoai/experimental/rfm/task_table.py,sha256=SPwkEdKRTwHHFcLqbvC1cDkyLXN2-3DpY5ujAyHRE-Q,8377
 kumoai/experimental/rfm/backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kumoai/experimental/rfm/backend/sqlite/__init__.py,sha256=jl-DBbhsqQ-dUXyWhyQTM1AU2qNAtXCmi1mokdhtBTg,902
-kumoai/experimental/rfm/backend/sqlite/table.py,sha256=PhBMj2vwChCxsTHrDzeiau8xs3EHgxA8t-uc6_JcREY,5221
-kumoai/experimental/rfm/backend/sqlite/sampler.py,sha256=pBCj0bXnh1PMg9cYJw-K983FzJX1_SLOe3QuOxkmvBQ,14242
+kumoai/experimental/rfm/backend/sqlite/table.py,sha256=WqYtd_rwlawItRMXZUfv14qdyU6huQmODuFjDo483dI,6683
+kumoai/experimental/rfm/backend/sqlite/sampler.py,sha256=_D9C5mj3oL4J2qZFap3emvTy2jxzth3dEWZPfr4dmEE,16201
 kumoai/experimental/rfm/backend/local/__init__.py,sha256=2s9sSA-E-8pfkkzCH4XPuaSxSznEURMfMgwEIfYYPsg,1014
-kumoai/experimental/rfm/backend/local/table.py,sha256=fCtLuBSvCSlTbrl77Pe1Aas4GQIC1TXfpuneBSv_wwI,3662
-kumoai/experimental/rfm/backend/local/graph_store.py,sha256=cY9KeLir9Xsp4MJl_K0VZckNa_LMKiiWCZG14uu21JI,11854
-kumoai/experimental/rfm/backend/local/sampler.py,sha256=85HoHCDiFOiuD_vFPZRx9JCyQUlLsqgsuB3NAw50wNw,10836
+kumoai/experimental/rfm/backend/local/table.py,sha256=GKeYGcu52ztCU8EBMqp5UVj85E145Ug41xiCPiTCXq4,3489
+kumoai/experimental/rfm/backend/local/graph_store.py,sha256=RHhkI13KpdPxqb4vXkwEwuFiX5DkrEsfZsOLywNnrvU,11294
+kumoai/experimental/rfm/backend/local/sampler.py,sha256=UKxTjsYs00sYuV_LAlDuZOvQq0BZzPCzZK1Fki2Fd70,10726
 kumoai/experimental/rfm/backend/snow/__init__.py,sha256=BYfsiuJ4Ee30GjG9EuUtitMHXnRfvVKi85zNlIwldV4,993
-kumoai/experimental/rfm/backend/snow/table.py,sha256=-LIumOr3jvd2Kww1gBUcFkZN_xaZEUGzhsLb2sRDMe8,5491
-kumoai/experimental/rfm/backend/snow/sampler.py,sha256=oNiBTo-Dr6LNCFJ9uHzLfKFYtx0rZq-Do2UOluwWWt0,10010
+kumoai/experimental/rfm/backend/snow/table.py,sha256=9N7TOcXX8hhAjCawnhuvQCArBFTCdng3gBakunUxg90,8892
+kumoai/experimental/rfm/backend/snow/sampler.py,sha256=zvPsgVnDfvskcnPWsIcqxw-Fn9DsCLfdoLE-m3bjeww,11483
 kumoai/experimental/rfm/pquery/__init__.py,sha256=X0O3EIq5SMfBEE-ii5Cq6iDhR3s3XMXB52Cx5htoePw,152
 kumoai/experimental/rfm/pquery/pandas_executor.py,sha256=MwSvFRwLq-z19LEdF0G0AT7Gj9tCqu-XLEA7mNbqXwc,18454
 kumoai/experimental/rfm/pquery/executor.py,sha256=gs5AVNaA50ci8zXOBD3qt5szdTReSwTs4BGuEyx4BEE,2728
-kumoai/experimental/rfm/infer/multicategorical.py,sha256=0-cLpDnGryhr76QhZNO-klKokJ6MUSfxXcGdQ61oykY,1102
+kumoai/experimental/rfm/infer/multicategorical.py,sha256=lNO_8aJw1whO6QVEMB3PRWMNlEEiX44g3v4tP88TSQY,1119
 kumoai/experimental/rfm/infer/categorical.py,sha256=VwNaKwKbRYkTxEJ1R6gziffC8dGsEThcDEfbi-KqW5c,853
 kumoai/experimental/rfm/infer/time_col.py,sha256=oNenUK6P7ql8uwShodtQ73uG1x3fbFWT78jRcF9DLTI,1789
 kumoai/experimental/rfm/infer/pkey.py,sha256=IaJI5GHK8ds_a3AOr3YYVgUlSmYYEgr4Nu92s2RyBV4,4412
 kumoai/experimental/rfm/infer/id.py,sha256=ZIO0DWIoiEoS_8MVc5lkqBfkTWWQ0yGCgjkwLdaYa_Q,908
-kumoai/experimental/rfm/infer/dtype.py,sha256=zjQ6Dhvb2QKv7A-c8mDxBnK5Dijhmn4fCgavIw4UH8Y,2582
-kumoai/experimental/rfm/infer/__init__.py,sha256=krdMFN8iKZlSFOl-M5MW1KuSviQV3H1E18jj2uB8g6Q,469
+kumoai/experimental/rfm/infer/dtype.py,sha256=FyAqvtrOWQC9hGrhQ7sC4BAI6c9k6ew-fo8ClS1sewM,2782
+kumoai/experimental/rfm/infer/__init__.py,sha256=8GDxQKd0pxZULdk7mpwl3CsOpL4v2HPuPEsbi2t_vzc,519
 kumoai/experimental/rfm/infer/timestamp.py,sha256=vM9--7eStzaGG13Y-oLYlpNJyhL6f9dp17HDXwtl_DM,1094
-kumoai/experimental/rfm/base/sql_sampler.py,sha256=-2dyftqvfbzMceIhE6i4wYFt7-p7FDeqlfH4P--qjWw,2598
-kumoai/experimental/rfm/base/__init__.py,sha256=ZWG3HDJfxlu0JZ6cdAbd9hmUCN1irLMrAXitMKmGlxs,799
-kumoai/experimental/rfm/base/sql_table.py,sha256=fX3aBvTgpt_S3XLGsUsWafeyV_SlO1Twah2ZbcI07GY,3764
-kumoai/experimental/rfm/base/table.py,sha256=5T52J_qNL8af84tZQZTD-9K4tDOtArWJRDH6jXOsdrQ,19888
-kumoai/experimental/rfm/base/column_expression.py,sha256=pYjM3nVsdBgsSFUlmx2nBfRj-0pch2ED6EfdS2WOUqk,336
-kumoai/experimental/rfm/base/sampler.py,sha256=aCD98t0CUhAvGXEFv24Vq2g4otuclpKkkyL1rMR_mFg,31449
-kumoai/experimental/rfm/base/source.py,sha256=RqlI_kBoRV0ADb8KdEKn15RNHMdFUzEVzb57lIoyBM4,294
-kumoai/experimental/rfm/base/column.py,sha256=Yc7oX8OUYJdnJmgD9Mz4jKSChjMXqJYRFybHWGVcDHE,2133
+kumoai/experimental/rfm/infer/stype.py,sha256=fu4zsOB-C7jNeMnq6dsK4bOZSewe7PtZe_AkohSRLoM,894
+kumoai/experimental/rfm/base/sql_sampler.py,sha256=qurkEVlMhDZw3d9SM2uGud6TMv_Wx_iqWoCgEKd_g9o,5094
+kumoai/experimental/rfm/base/__init__.py,sha256=rjmMux5lG8srw1bjQGcFQFv6zET9e5riP81nPkw28Jg,724
+kumoai/experimental/rfm/base/table.py,sha256=6qZeTMfnQejrn6TwqQeJGzJG7C0dSjJ7-NMLX38dvns,26563
+kumoai/experimental/rfm/base/sampler.py,sha256=tXYnVEyKC5NjSIpe8pNYp0V3Qbg-KbUE_QB0Emy2YiQ,30882
+kumoai/experimental/rfm/base/expression.py,sha256=Y7NtLTnKlx6euG_N3fLTcrFKheB6P5KS_jhCfoXV9DE,1252
+kumoai/experimental/rfm/base/source.py,sha256=bwu3GU2TvIXR2fwKAmJ1-5BDoNXMnI1SU3Fgdk8lWnc,301
+kumoai/experimental/rfm/base/column.py,sha256=GXzLC-VpShr6PecUzaj1MJKc_PHzfW5Jn9bOYPA8fFA,4965
 kumoai/encoder/__init__.py,sha256=VPGs4miBC_WfwWeOXeHhFomOUocERFavhKf5fqITcds,182
 kumoai/graph/graph.py,sha256=iyp4klPIMn2ttuEqMJvsrxKb_tmz_DTnvziIhCegduM,38291
 kumoai/graph/__init__.py,sha256=n8X4X8luox4hPBHTRC9R-3JzvYYMoR8n7lF1H4w4Hzc,228
@@ -56,6 +58,7 @@ kumoai/artifact_export/job.py,sha256=GEisSwvcjK_35RgOfsLXGgxMTXIWm765B_BW_Kgs-V0
 kumoai/artifact_export/__init__.py,sha256=BsfDrc3mCHpO9-BqvqKm8qrXDIwfdaoH5UIoG4eQkc4,238
 kumoai/utils/datasets.py,sha256=ptKIUoBONVD55pTVNdRCkQT3NWdN_r9UAUu4xewPa3U,2928
 kumoai/utils/__init__.py,sha256=6S-UtwjeLpnCYRCCIEWhkitPYGaqOGXC1ChE13DzXiU,256
+kumoai/utils/display.py,sha256=eXlw4B72y6zEruWYOfwvfqxfMBTL9AsPtWfw3BjaWqQ,1397
 kumoai/utils/progress_logger.py,sha256=3aYOoVSbQv5i9m2T8IqMydofKf6iNB1jxsl1uGjHZz8,9265
 kumoai/utils/sql.py,sha256=f6lR6rBEW7Dtk0NdM26dOZXUHDizEHb1WPlBCJrwoq0,118
 kumoai/utils/forecasting.py,sha256=-nDS6ucKNfQhTQOfebjefj0wwWH3-KYNslIomxwwMBM,7415
@@ -85,7 +88,7 @@ kumoai/connector/bigquery_connector.py,sha256=IkyRqvF8Cg96kApUuuz86eYnl-BqBmDX1f
 kumoai/connector/source_table.py,sha256=QLT8bEYaxeMwy-b168url0VfnkTrs5K6VKLbxTI4hEY,17539
 kumoai/connector/__init__.py,sha256=9g6oNJ0qHWFlL5enTSoK4_SSH_5hP74xUDZx-9SggC4,842
 kumoai/connector/file_upload_connector.py,sha256=swp03HgChOvmNPJetuujBSAqADe7NRmS_T0F3o9it4w,7008
-kumoai/connector/utils.py,sha256=wlqQxMmPvnFNoCcczGkKYjSu05h8OhWh4fhTzQm_2bQ,64694
+kumoai/connector/utils.py,sha256=sD3_Dmf42FobMfVayzMVkDHIfXzPN-htD3RHd6Kw8hQ,65055
 kumoai/connector/s3_connector.py,sha256=3kbv-h7DwD8O260Q0h1GPm5wwQpLt-Tb3d_CBSaie44,10155
 kumoai/connector/base.py,sha256=cujXSZF3zAfuxNuEw54DSL1T7XCuR4t0shSMDuPUagQ,5291
 kumoai/pquery/__init__.py,sha256=uTXr7t1eXcVfM-ETaM_1ImfEqhrmaj8BjiIvy1YZTL8,533
@@ -93,12 +96,12 @@ kumoai/pquery/predictive_query.py,sha256=UXn1s8ztubYZMNGl4ijaeidMiGlFveb1TGw9qI5
 kumoai/pquery/prediction_table.py,sha256=QPDH22X1UB0NIufY7qGuV2XW7brG3Pv--FbjNezzM2g,10776
 kumoai/pquery/training_table.py,sha256=elmPDZx11kPiC_dkOhJcBUGtHKgL32GCBvZ9k6U0pMg,15809
 kumoai/client/pquery.py,sha256=IQ8As-OOJOkuMoMosphOsA5hxQYLCbzOQJO7RezK8uY,7091
-kumoai/client/client.py,sha256=Jda8V9yiu3LbhxlcgRWPeYi7eF6jzCKcq8-B_vEd1ik,8514
+kumoai/client/client.py,sha256=npTLooBtmZ9xOo7AbEiYQTh9wFktsGSEpSEfdB7vdB4,8715
 kumoai/client/graph.py,sha256=zvLEDExLT_RVbUMHqVl0m6tO6s2gXmYSoWmPF6YMlnA,3831
 kumoai/client/online.py,sha256=pkBBh_DEC3GAnPcNw6bopNRlGe7EUbIFe7_seQqZRaw,2720
 kumoai/client/source_table.py,sha256=VCsCcM7KYcnjGP7HLTb-AOSEGEVsJTWjk8bMg1JdgPU,2101
 kumoai/client/__init__.py,sha256=MkyOuMaHQ2c8GPxjBDQSVFhfRE2d2_6CXQ6rxj4ps4w,64
-kumoai/client/jobs.py,sha256=iu_Wrta6BQMlV6ZtzSnmhjwNPKDMQDXOsqVVIyWodqw,17074
+kumoai/client/jobs.py,sha256=z3By5MWvWdJ_wYFyJA34pD4NueOXvXEqrAANWEpp4Pk,18066
 kumoai/client/utils.py,sha256=lz1NubwMDHCwzQRowRXm7mjAoYRd5UjRQIwXdtWAl90,3849
 kumoai/client/connector.py,sha256=x3i2aBTJTEMZvYRcWkY-UfWVOANZjqAso4GBbcshFjw,3920
 kumoai/client/table.py,sha256=cQG-RPm-e91idEgse1IPJDvBmzddIDGDkuyrR1rq4wU,3235
@@ -110,9 +113,10 @@ kumoai/trainer/job.py,sha256=Wk69nzFhbvuA3nEvtCstI04z5CxkgvQ6tHnGchE0Lkg,44938
 kumoai/trainer/baseline_trainer.py,sha256=LlfViNOmswNv4c6zJJLsyv0pC2mM2WKMGYx06ogtEVc,4024
 kumoai/trainer/__init__.py,sha256=zUdFl-f-sBWmm2x8R-rdVzPBeU2FaMzUY5mkcgoTa1k,939
 kumoai/trainer/online_serving.py,sha256=9cddb5paeZaCgbUeceQdAOxysCtV5XP-KcsgFz_XR5w,9566
+kumoai/trainer/distilled_trainer.py,sha256=2pPs5clakNxkLfaak7uqPJOrpTWe1RVVM7ztDSqQZvU,6484
 kumoai/trainer/trainer.py,sha256=hBXO7gwpo3t59zKFTeIkK65B8QRmWCwO33sbDuEAPlY,20133
-kumoai-2.14.0.dev202512181731.dist-info/RECORD,,
-kumoai-2.14.0.dev202512181731.dist-info/WHEEL,sha256=V1loQ6TpxABu1APUg0MoTRBOzSKT5xVc3skizX-ovCU,136
-kumoai-2.14.0.dev202512181731.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
-kumoai-2.14.0.dev202512181731.dist-info/METADATA,sha256=6DXWj2B67vLO-HDJaPEOX7fw4P5PqDzPDxRgf-o-UMY,2557
-kumoai-2.14.0.dev202512181731.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102
+kumoai-2.14.0.dev202601041732.dist-info/RECORD,,
+kumoai-2.14.0.dev202601041732.dist-info/WHEEL,sha256=V1loQ6TpxABu1APUg0MoTRBOzSKT5xVc3skizX-ovCU,136
+kumoai-2.14.0.dev202601041732.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
+kumoai-2.14.0.dev202601041732.dist-info/METADATA,sha256=dmLN7vMtkp6iM92XWX0BtKZCF8yU3RmoAKhJ-gwc4ME,2557
+kumoai-2.14.0.dev202601041732.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102

kumoai/experimental/rfm/base/column_expression.py

@@ -1,16 +0,0 @@
-from dataclasses import dataclass
-from typing import Any, TypeAlias
-
-from kumoapi.typing import Dtype
-
-from kumoai.mixin import CastMixin
-
-
-@dataclass(frozen=True)
-class ColumnExpressionSpec(CastMixin):
-    name: str
-    expr: str
-    dtype: Dtype | None = None
-
-
-ColumnExpressionType: TypeAlias = ColumnExpressionSpec | dict[str, Any]