oracle-ads 2.13.17__py3-none-any.whl → 2.13.18__py3-none-any.whl

ads/aqua/shaperecommend/shape_report.py

@@ -0,0 +1,233 @@
+#!/usr/bin/env python
+# Copyright (c) 2025 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+from ads.aqua.common.entities import ComputeShapeSummary
+from ads.aqua.shaperecommend.constants import QUANT_MAPPING
+from ads.aqua.shaperecommend.estimator import MemoryEstimator
+from ads.config import COMPARTMENT_OCID
+
+
+class RequestRecommend(BaseModel):
+    """
+    A request to recommend compute shapes and parameters for a given model.
+    """
+
+    model_id: str = Field(
+        ..., description="The OCID of the model to recommend feasible compute shapes."
+    )
+    generate_table: Optional[bool] = (
+        Field(
+            True,
+            description="True - to generate the rich diff Table, False - generate the JSON response",
+        ),
+    )
+    compartment_id: Optional[str] = Field(
+        COMPARTMENT_OCID, description="The OCID of user's compartment"
+    )
+
+    class Config:
+        protected_namespaces = ()
+
+
+class DeploymentParams(BaseModel):  # noqa: N801
+    """
+    Recommended parameters for deployment and model inferencing (specific to compute shape & model).
+    """
+
+    quantization: Optional[str] = Field(
+        None, description="Type of quantization (e.g. 4bit)."
+    )
+    max_model_len: int = Field(..., description="Maximum length of input sequence.")
+    params: str = Field(
+        ..., description="Runtime parameters for deployment with vLLM, etc."
+    )
+
+
+class ModelDetail(BaseModel):
+    """
+    The estimated memory footprint of a model, KV cache, and its total (model + KV cache).
+    """
+
+    model_size_gb: float = Field(..., description="Size of the model in GB.")
+    kv_cache_size_gb: float = Field(..., description="Size of KV cache in GB.")
+    total_model_gb: float = Field(
+        ..., description="Total size of model and cache in GB."
+    )
+
+    class Config:
+        protected_namespaces = ()
+
+
+class ModelConfig(BaseModel):
+    """
+    The configuration for a model based on specific set of deployment parameters and memory capacity of shape.
+    """
+
+    model_details: ModelDetail = Field(..., description="Details about the model.")
+    deployment_params: DeploymentParams = Field(
+        ..., description="Parameters for deployment."
+    )
+    recommendation: str = Field(..., description="GPU recommendation for the model.")
+
+    class Config:
+        protected_namespaces = ()
+
+    @classmethod
+    def constuct_model_config(
+        cls, estimator: MemoryEstimator, allowed_gpu_memory: float
+    ) -> "ModelConfig":
+        """
+        Assembles a complete ModelConfig, including model details, deployment parameters (vLLM), and recommendations.
+
+        Parameters
+        ----------
+        shape_quantization : set[str]
+            Allowed quantization methods for the compute shape
+
+        Returns
+        -------
+        ModelConfig
+            Contains round-tripped model size, kv cache, total, vLLM parameters, and recommendations.
+
+        Notes
+        -----
+        - Rounds all sizes to 3 decimal digits.
+        - Computes a recommendation string using `limiting_factor`.
+        """
+        c = estimator.llm_config
+        deployment_params = DeploymentParams(
+            quantization=c.quantization or c.in_flight_quantization or c.weight_dtype,
+            max_model_len=getattr(estimator, "seq_len", None),
+            params=estimator.construct_deployment_params(),
+        )
+        model_detail = ModelDetail(
+            model_size_gb=round(getattr(estimator, "model_memory", 0.0), 2),
+            kv_cache_size_gb=round(getattr(estimator, "kv_cache_memory", 0.0), 2),
+            total_model_gb=round(getattr(estimator, "total_memory", 0.0), 2),
+        )
+        return ModelConfig(
+            model_details=model_detail,
+            deployment_params=deployment_params,
+            recommendation=estimator.limiting_factor(allowed_gpu_memory),
+        )
+
+
+class ShapeReport(BaseModel):
+    """
+    The feasible deployment configurations for the model per shape.
+    """
+
+    shape_details: "ComputeShapeSummary" = Field(
+        ..., description="Details about the compute shape (ex. VM.GPU.A10.2)."
+    )
+    configurations: List["ModelConfig"] = Field(
+        default_factory=list, description="List of model configurations."
+    )
+
+    def is_dominated(self, others: List["ShapeReport"]) -> bool:
+        """
+        Determines whether this shape is dominated by any other shape in a Pareto sense.
+
+        Parameters
+        ----------
+        others : list of ShapeReport
+            List of other shape/deployment configurations to compare against.
+
+        Returns
+        -------
+        bool
+            True if this shape is dominated by at least one other, False otherwise.
+
+        Notes
+        -----
+        A shape is dominated if there exists another configuration that is
+        at least as good in all criteria and strictly better in at least one.
+        Criteria:
+        - Cost (to be minimized)
+        - Performance, quantization level, max sequence length (to be maximized)
+        """
+        try:
+            cand_cost = self.shape_details.gpu_specs.ranking.cost
+            cand_perf = self.shape_details.gpu_specs.ranking.performance
+            cand_quant = QUANT_MAPPING.get(
+                self.configurations[0].deployment_params.quantization, 0
+            )
+            cand_maxlen = self.configurations[0].deployment_params.max_model_len
+
+            for other in others:
+                other_cost = other.shape_details.gpu_specs.ranking.cost
+                other_perf = other.shape_details.gpu_specs.ranking.performance
+                other_quant = QUANT_MAPPING.get(
+                    other.configurations[0].deployment_params.quantization, 0
+                )
+                other_maxlen = other.configurations[0].deployment_params.max_model_len
+                if (
+                    other_cost <= cand_cost
+                    and other_perf >= cand_perf
+                    and other_quant >= cand_quant
+                    and other_maxlen >= cand_maxlen
+                    and (
+                        other_cost < cand_cost
+                        or other_perf > cand_perf
+                        or other_quant > cand_quant
+                        or other_maxlen > cand_maxlen
+                    )
+                ):
+                    return True
+            return False
+        except AttributeError:
+            return False
+
+    @classmethod
+    def pareto_front(cls, shapes: List["ShapeReport"]) -> List["ShapeReport"]:
+        """
+        Filters a list of shapes/configurations to those on the Pareto frontier.
+
+        Parameters
+        ----------
+        shapes : list of ShapeReport
+            List of candidate shape/configuration reports to evaluate.
+
+        Returns
+        -------
+        list of ShapeReport
+            Subset of input shapes that are not dominated by any other (the Pareto front).
+
+        Notes
+        -----
+        The returned set contains non-dominated deployments for maximizing
+        performance, quantization, and model length, while minimizing cost.
+        """
+        return [
+            shape
+            for shape in shapes
+            if not shape.is_dominated([s for s in shapes if s != shape])
+        ]
+
+
+class ShapeRecommendationReport(BaseModel):
+    """
+    Full report of shape fit recommendations and troubleshooting, if applicable.
+
+    Attributes:
+        recommendations (List[DeploymentShapeSummary]): Recommended deployment shapes
+            for each tested batch size and max sequence length combination.
+        troubleshoot (Optional[TroubleshootShapeSummary]): Troubleshooting information
+            if no valid deployment shapes are available.
+    """
+
+    display_name: Optional[str] = Field(
+        "", description="Name of the model used for recommendations."
+    )
+    recommendations: List[ShapeReport] = Field(
+        default_factory=list, description="List of shape fit recommendations."
+    )
+    troubleshoot: Optional[str] = Field(
+        None,
+        description="Details for troubleshooting if no shapes fit the current model.",
+    )

ads/aqua/version.json

@@ -1,3 +1,3 @@
 {
-  "aqua": "1.0.7a"
+  "aqua": "1.0.8"
 }

ads/cli.py

@@ -7,6 +7,8 @@ import logging
 import sys
 import traceback
 import uuid
+from rich.console import Console
+from rich.table import Table
 
 import fire
 from pydantic import BaseModel
@@ -92,6 +94,12 @@ def serialize(data):
                 print(str(item))
     elif isinstance(data, BaseModel):
         print(json.dumps(data.dict(), indent=4))
+    elif isinstance(data, Table):
+        console = Console()
+        console.print(data)
+        return
+    elif data is None:
+        return
     else:
         print(str(data))
 
@@ -131,7 +139,7 @@ def exit_program(ex: Exception, logger: "logging.Logger") -> None:
 
     request_id = str(uuid.uuid4())
     logger.debug(f"Error Request ID: {request_id}\nError: {traceback.format_exc()}")
-    logger.error(f"Error Request ID: {request_id}\n" f"Error: {str(ex)}")
+    logger.error(f"Error Request ID: {request_id}\nError: {str(ex)}")
 
     exit_code = getattr(ex, "exit_code", 1)
     logger.error(f"Exit code: {exit_code}")

ads/jobs/builders/infrastructure/dsc_job.py

@@ -1751,6 +1751,7 @@ class DataScienceJob(Infrastructure):
         return (
             MULTI_NODE_JOB_SUPPORT
             and isinstance(runtime, MultiNodeRuntime)
+            and runtime.replica
             and runtime.replica > 1
         )
 

ads/jobs/builders/infrastructure/dsc_job_runtime.py

@@ -365,6 +365,11 @@ class RuntimeHandler:
             dsc_job,
             "job_node_configuration_details.job_node_group_configuration_details_list",
         )
+        if node_groups is None:
+            node_groups = get_value(
+                dsc_job,
+                "job_node_configuration_details.jobNodeGroupConfigurationDetailsList",
+            )
         if node_groups and len(node_groups) == 1:
             return node_groups[0]
         return None
@@ -373,6 +378,7 @@ class RuntimeHandler:
         node_group = self._get_node_group(dsc_job)
         if node_group:
             replica = get_value(node_group, "replicas")
+            envs.pop(self.CONST_NODE_COUNT, None)
         elif not envs:
             replica = None
         elif self.CONST_WORKER_COUNT in envs:
@@ -399,7 +405,9 @@ class RuntimeHandler:
         env_attr = "job_configuration_details.environment_variables"
         node_group = self._get_node_group(dsc_job)
         if node_group:
-            envs = get_value(node_group, env_attr)
+            envs = get_value(node_group, env_attr) or get_value(
+                node_group, "jobConfigurationDetails.environment_variables"
+            )
         else:
             envs = get_value(dsc_job, env_attr)
         if envs:

ads/model/service/oci_datascience_model_deployment.py

@@ -1,23 +1,24 @@
 #!/usr/bin/env python
-# -*- coding: utf-8; -*-
 
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
-from functools import wraps
 import logging
-from typing import Callable, List
-from ads.common.oci_datascience import OCIDataScienceMixin
-from ads.common.work_request import DataScienceWorkRequest
-from ads.config import PROJECT_OCID
-from ads.model.deployment.common.utils import OCIClientManager, State
-import oci
+from functools import wraps
+from typing import Callable, List, Optional
 
+import oci
 from oci.data_science.models import (
     CreateModelDeploymentDetails,
+    ModelDeploymentShapeSummary,
     UpdateModelDeploymentDetails,
 )
 
+from ads.common.oci_datascience import OCIDataScienceMixin
+from ads.common.work_request import DataScienceWorkRequest
+from ads.config import COMPARTMENT_OCID, PROJECT_OCID
+from ads.model.deployment.common.utils import OCIClientManager, State
+
 DEFAULT_WAIT_TIME = 1200
 DEFAULT_POLL_INTERVAL = 10
 ALLOWED_STATUS = [
@@ -185,14 +186,13 @@ class OCIDataScienceModelDeployment(
                 self.id,
             )
 
-           
             self.workflow_req_id = response.headers.get("opc-work-request-id", None)
             if wait_for_completion:
                 try:
                     DataScienceWorkRequest(self.workflow_req_id).wait_work_request(
                         progress_bar_description="Activating model deployment",
-                        max_wait_time=max_wait_time, 
-                        poll_interval=poll_interval
+                        max_wait_time=max_wait_time,
+                        poll_interval=poll_interval,
                     )
                 except Exception as e:
                     logger.error(
@@ -239,8 +239,8 @@ class OCIDataScienceModelDeployment(
             try:
                 DataScienceWorkRequest(self.workflow_req_id).wait_work_request(
                     progress_bar_description="Creating model deployment",
-                    max_wait_time=max_wait_time, 
-                    poll_interval=poll_interval
+                    max_wait_time=max_wait_time,
+                    poll_interval=poll_interval,
                 )
             except Exception as e:
                 logger.error("Error while trying to create model deployment: " + str(e))
@@ -290,8 +290,8 @@ class OCIDataScienceModelDeployment(
                 try:
                     DataScienceWorkRequest(self.workflow_req_id).wait_work_request(
                         progress_bar_description="Deactivating model deployment",
-                        max_wait_time=max_wait_time, 
-                        poll_interval=poll_interval
+                        max_wait_time=max_wait_time,
+                        poll_interval=poll_interval,
                     )
                 except Exception as e:
                     logger.error(
@@ -351,14 +351,14 @@ class OCIDataScienceModelDeployment(
         response = self.client.delete_model_deployment(
             self.id,
         )
-        
+
         self.workflow_req_id = response.headers.get("opc-work-request-id", None)
         if wait_for_completion:
             try:
                 DataScienceWorkRequest(self.workflow_req_id).wait_work_request(
                     progress_bar_description="Deleting model deployment",
-                    max_wait_time=max_wait_time, 
-                    poll_interval=poll_interval
+                    max_wait_time=max_wait_time,
+                    poll_interval=poll_interval,
                 )
             except Exception as e:
                 logger.error("Error while trying to delete model deployment: " + str(e))
@@ -493,3 +493,30 @@ class OCIDataScienceModelDeployment(
             An instance of `OCIDataScienceModelDeployment`.
         """
         return super().from_ocid(model_deployment_id)
+
+    @classmethod
+    def shapes(
+        cls,
+        compartment_id: Optional[str] = None,
+        **kwargs,
+    ) -> List[ModelDeploymentShapeSummary]:
+        """
+        Retrieves all available model deployment shapes in the given compartment.
+
+        This method uses OCI's pagination utility to fetch all pages of model
+        deployment shape summaries available in the specified compartment.
+
+        Args:
+            compartment_id (Optional[str]): The OCID of the compartment. If not provided,
+                                            the default COMPARTMENT_ID extracted form env variables is used.
+            **kwargs: Additional keyword arguments to pass to the list_model_deployments call.
+
+        Returns:
+            List[ModelDeploymentShapeSummary]: A list of all model deployment shape summaries.
+        """
+        client = cls().client
+        compartment_id = compartment_id or COMPARTMENT_OCID
+
+        return oci.pagination.list_call_get_all_results(
+            client.list_model_deployment_shapes, compartment_id, **kwargs
+        ).data

ads/opctl/operator/lowcode/common/data.py

@@ -19,16 +19,21 @@ from .transformations import Transformations
 
 
 class AbstractData(ABC):
-    def __init__(self, spec, name="input_data", data=None):
+    def __init__(self, spec, name="input_data", data=None, subset=None):
         self.Transformations = Transformations
         self.data = None
         self._data_dict = dict()
         self.name = name
         self.spec = spec
+        self.subset = subset
         if data is not None:
             self.data = data
         else:
             self.load_transform_ingest_data(spec)
+        # Subset by series if requested
+        # if self.subset is not None and hasattr(self, 'data') and self.data is not None:
+        #     subset_str = [str(s) for s in self.subset]
+        #     self.data = self.data[self.data.index.get_level_values(DataColumns.Series).isin(subset_str)]
 
     def get_raw_data_by_cat(self, category):
         mapping = self._data_transformer.get_target_category_columns_map()
@@ -72,7 +77,7 @@ class AbstractData(ABC):
     def _load_data(self, data_spec, **kwargs):
         loading_start_time = time.time()
         try:
-            raw_data = load_data(data_spec)
+            raw_data = load_data(data_spec, subset=self.subset if self.subset else None, target_category_columns=self.spec.target_category_columns)
         except InvalidParameterError as e:
             e.args = e.args + (f"Invalid Parameter: {self.name}",)
             raise e

ads/opctl/operator/lowcode/common/transformations.py

@@ -294,3 +294,210 @@ class Transformations(ABC):
     def _fill_na(self, df: pd.DataFrame, na_value=0) -> pd.DataFrame:
         """Fill nans in dataframe"""
         return df.fillna(value=na_value)
+
+    def build_fforms_meta_features(self, data, target_col=None, group_cols=None):
+        """
+        Build meta-features for time series based on FFORMS paper and add them to the original DataFrame.
+
+        Parameters
+        ----------
+        data : pandas.DataFrame
+            Input DataFrame containing time series data
+        target_col : str, optional
+            Name of the target column to calculate meta-features for.
+            If None, uses the target column specified in dataset_info.
+        group_cols : list of str, optional
+            List of columns to group by before calculating meta-features.
+            If None, calculates features for the entire series.
+
+        Returns
+        -------
+        pandas.DataFrame
+            Original DataFrame with additional meta-feature columns
+
+        References
+        ----------
+        Talagala, T. S., Hyndman, R. J., & Athanasopoulos, G. (2023).
+        Meta-learning how to forecast time series. Journal of Forecasting, 42(6), 1476-1501.
+        """
+        if not isinstance(data, pd.DataFrame):
+            raise ValueError("Input must be a pandas DataFrame")
+
+        # Use target column from dataset_info if not specified
+        if target_col is None:
+            target_col = self.target_column_name
+        if target_col not in data.columns:
+            raise ValueError(f"Target column '{target_col}' not found in DataFrame")
+
+        # Check if group_cols are provided and valid
+        if group_cols is not None:
+            if not isinstance(group_cols, list):
+                raise ValueError("group_cols must be a list of column names")
+            for col in group_cols:
+                if col not in data.columns:
+                    raise ValueError(f"Group column '{col}' not found in DataFrame")
+
+        # If no group_cols, get the target_category_columns else treat the entire DataFrame as a single series
+        if not group_cols:
+            group_cols = self.target_category_columns if self.target_category_columns else []
+
+        # Calculate meta-features for each series
+        def calculate_series_features(series):
+            """Calculate features for a single series"""
+            n = len(series)
+            values = series.values
+
+            # Basic statistics
+            mean = series.mean()
+            std = series.std()
+            variance = series.var()
+            skewness = series.skew()
+            kurtosis = series.kurtosis()
+            cv = std / mean if mean != 0 else np.inf
+
+            # Trend features
+            X = np.vstack([np.arange(n), np.ones(n)]).T
+            trend_coef = np.linalg.lstsq(X, values, rcond=None)[0][0]
+            trend_pred = X.dot(np.linalg.lstsq(X, values, rcond=None)[0])
+            residuals = values - trend_pred
+            std_residuals = np.std(residuals)
+
+            # Turning points
+            turning_points = 0
+            for i in range(1, n-1):
+                if (values[i-1] < values[i] and values[i] > values[i+1]) or \
+                   (values[i-1] > values[i] and values[i] < values[i+1]):
+                    turning_points += 1
+            turning_points_rate = turning_points / (n-2) if n > 2 else 0
+
+            # Serial correlation
+            acf1 = series.autocorr(lag=1) if n > 1 else 0
+            acf2 = series.autocorr(lag=2) if n > 2 else 0
+            acf10 = series.autocorr(lag=10) if n > 10 else 0
+
+            # Seasonality features
+            seasonal_strength = 0
+            seasonal_peak_strength = 0
+            if n >= 12:
+                seasonal_lags = [12, 24, 36]
+                seasonal_acfs = []
+                for lag in seasonal_lags:
+                    if n > lag:
+                        acf_val = series.autocorr(lag=lag)
+                        seasonal_acfs.append(abs(acf_val))
+                seasonal_peak_strength = max(seasonal_acfs) if seasonal_acfs else 0
+
+                ma = series.rolling(window=12, center=True).mean()
+                seasonal_comp = series - ma
+                seasonal_strength = 1 - np.var(seasonal_comp.dropna()) / np.var(series)
+
+            # Stability and volatility features
+            values_above_mean = values >= mean
+            crossing_points = np.sum(values_above_mean[1:] != values_above_mean[:-1])
+            crossing_rate = crossing_points / (n - 1) if n > 1 else 0
+
+            # First and second differences
+            diff1 = np.diff(values)
+            diff2 = np.diff(diff1) if len(diff1) > 1 else np.array([])
+
+            diff1_mean = np.mean(np.abs(diff1)) if len(diff1) > 0 else 0
+            diff1_var = np.var(diff1) if len(diff1) > 0 else 0
+            diff2_mean = np.mean(np.abs(diff2)) if len(diff2) > 0 else 0
+            diff2_var = np.var(diff2) if len(diff2) > 0 else 0
+
+            # Nonlinearity features
+            if n > 3:
+                X = values[:-1].reshape(-1, 1)
+                y = values[1:]
+                X2 = X * X
+                X3 = X * X * X
+                X_aug = np.hstack([X, X2, X3])
+                nonlinearity = np.linalg.lstsq(X_aug, y, rcond=None)[1][0] if len(y) > 0 else 0
+            else:
+                nonlinearity = 0
+
+            # Long-term trend features
+            if n >= 10:
+                mid = n // 2
+                trend_change = np.mean(values[mid:]) - np.mean(values[:mid])
+            else:
+                trend_change = 0
+
+            # Step changes and spikes
+            step_changes = np.abs(diff1).max() if len(diff1) > 0 else 0
+            spikes = np.sum(np.abs(values - mean) > 2 * std) / n if std != 0 else 0
+
+            # Hurst exponent and entropy
+            lag = min(10, n // 2)
+            variance_ratio = np.var(series.diff(lag)) / (lag * np.var(series.diff())) if n > lag else 0
+            hurst = np.log(variance_ratio) / (2 * np.log(lag)) if variance_ratio > 0 and lag > 1 else 0
+
+            hist, _ = np.histogram(series, bins='auto', density=True)
+            entropy = -np.sum(hist[hist > 0] * np.log(hist[hist > 0]))
+
+            return pd.Series({
+                'ts_n_obs': n,
+                'ts_mean': mean,
+                'ts_std': std,
+                'ts_variance': variance,
+                'ts_cv': cv,
+                'ts_skewness': skewness,
+                'ts_kurtosis': kurtosis,
+                'ts_trend': trend_coef,
+                'ts_trend_change': trend_change,
+                'ts_std_residuals': std_residuals,
+                'ts_turning_points_rate': turning_points_rate,
+                'ts_seasonal_strength': seasonal_strength,
+                'ts_seasonal_peak_strength': seasonal_peak_strength,
+                'ts_acf1': acf1,
+                'ts_acf2': acf2,
+                'ts_acf10': acf10,
+                'ts_crossing_rate': crossing_rate,
+                'ts_diff1_mean': diff1_mean,
+                'ts_diff1_variance': diff1_var,
+                'ts_diff2_mean': diff2_mean,
+                'ts_diff2_variance': diff2_var,
+                'ts_nonlinearity': nonlinearity,
+                'ts_step_max': step_changes,
+                'ts_spikes_rate': spikes,
+                'ts_hurst': hurst,
+                'ts_entropy': entropy
+            })
+
+        # Create copy of input DataFrame
+        result_df = data.copy()
+
+        if group_cols:
+            # Calculate features for each group
+            features = []
+            # Sort by date within each group if date column exists
+            date_col = self.dt_column_name if self.dt_column_name else 'Date'
+            if date_col in data.columns:
+                data = data.sort_values([date_col] + group_cols)
+
+            for name, group in data.groupby(group_cols):
+                # Sort group by date if exists
+                if date_col in group.columns:
+                    group = group.sort_values(date_col)
+                group_features = calculate_series_features(group[target_col])
+                if isinstance(name, tuple):
+                    feature_row = dict(zip(group_cols, name))
+                else:
+                    feature_row = {group_cols[0]: name}
+                feature_row.update(group_features)
+                features.append(feature_row)
+
+            # Create features DataFrame without merging
+            features_df = pd.DataFrame(features)
+            # Return only the meta-features DataFrame with group columns
+            return features_df
+        else:
+            # Sort by date if exists and calculate features for entire series
+            date_col = self.dt_column_name if self.dt_column_name else 'Date'
+            if date_col in data.columns:
+                data = data.sort_values(date_col)
+            features = calculate_series_features(data[target_col])
+            # Return single row DataFrame with meta-features
+            return pd.DataFrame([features])
+
+        return result_df

ads/opctl/operator/lowcode/common/utils.py

@@ -124,6 +124,14 @@ def load_data(data_spec, storage_options=None, **kwargs):
         data = data[columns]
     if limit:
         data = data[:limit]
+    # Filtering by subset if provided
+    subset = kwargs.get('subset', None)
+    if subset is not None:
+        target_category_columns = kwargs.get('target_category_columns', None)
+        mask = False
+        for col in target_category_columns:
+            mask = mask | data[col].isin(subset)
+            data = data[mask]
     return data