nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/data_models/telemetry_exporter.py

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+from nat.data_models.common import BaseModelRegistryTag
+from nat.data_models.common import TypedBaseModel
+
+
+class TelemetryExporterBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+TelemetryExporterConfigT = typing.TypeVar("TelemetryExporterConfigT", bound=TelemetryExporterBaseConfig)

nat/data_models/temperature_mixin.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+from pydantic import BaseModel
+
+from nat.data_models.gated_field_mixin import GatedFieldMixin
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import SearchSpace
+
+
+class TemperatureMixin(
+        BaseModel,
+        GatedFieldMixin,
+        field_name="temperature",
+        default_if_supported=0.0,
+        keys=("model_name", "model", "azure_deployment"),
+        unsupported=(re.compile(r"gpt-?5", re.IGNORECASE), ),
+):
+    """
+    Mixin class for temperature configuration. Unsupported on models like gpt-5.
+
+    Attributes:
+        temperature: Sampling temperature in [0, 1]. Defaults to 0.0 when supported on the model.
+    """
+    temperature: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        le=1.0,
+        description="Sampling temperature in [0, 1]. Defaults to 0.0 when supported on the model.",
+        space=SearchSpace(high=0.9, low=0.1, step=0.2))

nat/data_models/thinking_mixin.py

@@ -0,0 +1,86 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.data_models.gated_field_mixin import GatedFieldMixin
+
+# Currently the control logic for thinking is only implemented for Nemotron models
+_NEMOTRON_REGEX = re.compile(r"^nvidia/(llama|nvidia).*nemotron", re.IGNORECASE)
+# The keys are the fields that are used to determine if the model supports thinking
+_MODEL_KEYS = ("model_name", "model", "azure_deployment")
+
+
+class ThinkingMixin(
+        BaseModel,
+        GatedFieldMixin,
+        field_name="thinking",
+        default_if_supported=None,
+        keys=_MODEL_KEYS,
+        supported=(_NEMOTRON_REGEX, ),
+):
+    """
+    Mixin class for thinking configuration. Only supported on Nemotron models.
+
+    Attributes:
+        thinking: Whether to enable thinking. Defaults to None when supported on the model.
+    """
+    thinking: bool | None = Field(
+        default=None,
+        description="Whether to enable thinking. Defaults to None when supported on the model.",
+    )
+
+    @property
+    def thinking_system_prompt(self) -> str | None:
+        """
+        Returns the system prompt to use for thinking.
+        For NVIDIA Nemotron, returns "/think" if enabled, else "/no_think".
+        For Llama Nemotron v1.5, returns "/think" if enabled, else "/no_think".
+        For Llama Nemotron v1.0 or v1.1, returns "detailed thinking on" if enabled, else "detailed thinking off".
+        If thinking is not supported on the model, returns None.
+
+        Returns:
+            str | None: The system prompt to use for thinking.
+        """
+        if self.thinking is None:
+            return None
+
+        for key in _MODEL_KEYS:
+            model = getattr(self, key, None)
+            if not isinstance(model, str) or model is None:
+                continue
+
+            # Normalize name to reduce checks
+            model = model.lower().translate(str.maketrans("_.", "--"))
+
+            if model.startswith("nvidia/nvidia"):
+                return "/think" if self.thinking else "/no_think"
+
+            if model.startswith("nvidia/llama"):
+                if "v1-0" in model or "v1-1" in model or model.endswith("v1"):
+                    return f"detailed thinking {'on' if self.thinking else 'off'}"
+
+                if "v1-5" in model:
+                    # v1.5 models are updated to use the /think and /no_think system prompts
+                    return "/think" if self.thinking else "/no_think"
+
+                # Assume any other model is a newer model that uses the /think and /no_think system prompts
+                return "/think" if self.thinking else "/no_think"
+
+        # Unknown model
+        return None

nat/data_models/top_p_mixin.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+from pydantic import BaseModel
+
+from nat.data_models.gated_field_mixin import GatedFieldMixin
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import SearchSpace
+
+
+class TopPMixin(
+        BaseModel,
+        GatedFieldMixin,
+        field_name="top_p",
+        default_if_supported=1.0,
+        keys=("model_name", "model", "azure_deployment"),
+        unsupported=(re.compile(r"gpt-?5", re.IGNORECASE), ),
+):
+    """
+    Mixin class for top-p configuration. Unsupported on models like gpt-5.
+
+    Attributes:
+        top_p: Top-p for distribution sampling. Defaults to 1.0 when supported on the model.
+    """
+    top_p: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        le=1.0,
+        description="Top-p for distribution sampling. Defaults to 1.0 when supported on the model.",
+        space=SearchSpace(high=1.0, low=0.5, step=0.1))

nat/data_models/ttc_strategy.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class TTCStrategyBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    Base configuration class for Test Time Compute (TTC) strategy.
+    This class is used to define the structure of TTC strategy configurations.
+    """
+    pass
+
+
+TTCStrategyBaseConfigT = typing.TypeVar("TTCStrategyBaseConfigT", bound=TTCStrategyBaseConfig)

nat/embedder/azure_openai_embedder.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.embedder import EmbedderProviderInfo
+from nat.cli.register_workflow import register_embedder_provider
+from nat.data_models.embedder import EmbedderBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
+
+
+class AzureOpenAIEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="azure_openai"):
+    """An Azure OpenAI embedder provider to be used with an embedder client."""
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    api_key: str | None = Field(default=None, description="Azure OpenAI API key to interact with hosted model.")
+    api_version: str = Field(default="2025-04-01-preview", description="Azure OpenAI API version.")
+    azure_endpoint: str | None = Field(validation_alias=AliasChoices("azure_endpoint", "base_url"),
+                                       serialization_alias="azure_endpoint",
+                                       default=None,
+                                       description="Base URL for the hosted model.")
+    azure_deployment: str = Field(validation_alias=AliasChoices("azure_deployment", "model_name", "model"),
+                                  serialization_alias="azure_deployment",
+                                  description="The Azure OpenAI hosted model/deployment name.")
+
+
+@register_embedder_provider(config_type=AzureOpenAIEmbedderModelConfig)
+async def azure_openai_embedder_model(config: AzureOpenAIEmbedderModelConfig, _builder: Builder):
+
+    yield EmbedderProviderInfo(config=config, description="An Azure OpenAI model for use with an Embedder client.")

nat/embedder/nim_embedder.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+from pydantic import AfterValidator
+from pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.embedder import EmbedderProviderInfo
+from nat.cli.register_workflow import register_embedder_provider
+from nat.data_models.embedder import EmbedderBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
+
+allowed_truncate_values = ["NONE", "START", "END"]
+
+
+def option_in_allowed_values(v):
+    """Ensures option is allowed"""
+    assert v in allowed_truncate_values
+    return v
+
+
+TruncationOption = typing.Annotated[str, AfterValidator(option_in_allowed_values)]
+
+
+class NIMEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="nim"):
+    """A NVIDIA Inference Microservice (NIM) embedder provider to be used with an embedder client."""
+
+    api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")
+    base_url: str | None = Field(default=None, description="Base url to the hosted NIM.")
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The model name for the hosted NIM.")
+    truncate: TruncationOption = Field(default="NONE",
+                                       description=("The truncation strategy if the input on the "
+                                                    "server side if it's too large."))
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+
+@register_embedder_provider(config_type=NIMEmbedderModelConfig)
+async def nim_embedder_model(embedder_config: NIMEmbedderModelConfig, builder: Builder):
+
+    yield EmbedderProviderInfo(config=embedder_config, description="A NIM model for use with an Embedder client.")

nat/embedder/openai_embedder.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.embedder import EmbedderProviderInfo
+from nat.cli.register_workflow import register_embedder_provider
+from nat.data_models.embedder import EmbedderBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
+
+
+class OpenAIEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="openai"):
+    """An OpenAI LLM provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    api_key: str | None = Field(default=None, description="OpenAI API key to interact with hosted model.")
+    base_url: str | None = Field(default=None, description="Base url to the hosted model.")
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The OpenAI hosted model name.")
+
+
+@register_embedder_provider(config_type=OpenAIEmbedderModelConfig)
+async def openai_embedder_model(config: OpenAIEmbedderModelConfig, _builder: Builder):
+
+    yield EmbedderProviderInfo(config=config, description="An OpenAI model for use with an Embedder client.")

nat/embedder/register.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# flake8: noqa
+# isort:skip_file
+
+# Import any providers which need to be automatically registered here
+from . import azure_openai_embedder
+from . import nim_embedder
+from . import openai_embedder

nat/eval/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/eval/config.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from nat.eval.evaluator.evaluator_model import EvalInput
+from nat.eval.evaluator.evaluator_model import EvalOutput
+from nat.eval.usage_stats import UsageStats
+from nat.profiler.data_models import ProfilerResults
+
+
+class EvaluationRunConfig(BaseModel):
+    """
+    Parameters used for a single evaluation run.
+    """
+    config_file: Path | BaseModel
+    dataset: str | None = None  # dataset file path can be specified in the config file
+    result_json_path: str = "$"
+    skip_workflow: bool = False
+    skip_completed_entries: bool = False
+    endpoint: str | None = None  # only used when running the workflow remotely
+    endpoint_timeout: int = 300
+    reps: int = 1
+    override: tuple[tuple[str, str], ...] = ()
+    # If false, the output will not be written to the output directory. This is
+    # useful when running evaluation via another tool.
+    write_output: bool = True
+    # if true, the dataset is adjusted to a multiple of the concurrency
+    adjust_dataset_size: bool = False
+    # number of passes at each concurrency, if 0 the dataset is adjusted to a multiple of the
+    # concurrency. The is only used if adjust_dataset_size is true
+    num_passes: int = 0
+    # timeout for waiting for trace export tasks to complete
+    export_timeout: float = 60.0
+
+
+class EvaluationRunOutput(BaseModel):
+    """
+    Output of a single evaluation run.
+    """
+    workflow_output_file: Path | None
+    evaluator_output_files: list[Path]
+    workflow_interrupted: bool
+
+    eval_input: EvalInput
+    evaluation_results: list[tuple[str, EvalOutput]]
+    usage_stats: UsageStats | None = None
+    profiler_results: ProfilerResults

nat/eval/dataset_handler/dataset_downloader.py

@@ -0,0 +1,106 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+from pathlib import Path
+
+import boto3
+import requests
+from botocore.exceptions import NoCredentialsError
+
+from nat.data_models.dataset_handler import EvalDatasetConfig
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetDownloader:
+    """
+    Download remote datasets using signed URLs or S3 credentials.
+
+    One DatasetDownloader object is needed for each dataset to be downloaded.
+    """
+
+    def __init__(self, dataset_config: EvalDatasetConfig):
+        self.dataset_config = dataset_config
+        self._s3_client = None
+
+    @property
+    def s3_config(self):
+        return self.dataset_config.s3
+
+    @property
+    def s3_client(self):
+        """Lazy init the S3 client."""
+        if not self._s3_client:
+            try:
+                self._s3_client = boto3.client("s3",
+                                               endpoint_url=self.s3_config.endpoint_url,
+                                               aws_access_key_id=self.s3_config.access_key,
+                                               aws_secret_access_key=self.s3_config.secret_key)
+            except NoCredentialsError as e:
+                logger.error("AWS credentials not available: %s", e)
+                raise
+            except Exception as e:
+                logger.error("Failed to initialize S3 client: %s", e)
+                raise
+        return self._s3_client
+
+    @staticmethod
+    def ensure_directory_exists(file_path: str):
+        """Ensure the directory for the file exists."""
+        Path(file_path).parent.mkdir(parents=True, exist_ok=True)
+
+    def download_with_signed_url(self, remote_file_path: str, local_file_path: str, timeout: int = 300):
+        """Download a file using a signed URL."""
+        try:
+            response = requests.get(remote_file_path, stream=True, timeout=timeout)
+            response.raise_for_status()
+            with open(local_file_path, "wb") as file:
+                for chunk in response.iter_content(chunk_size=8192):
+                    file.write(chunk)
+            logger.info("File downloaded successfully to %s using signed URL.", local_file_path)
+        except requests.exceptions.RequestException as e:
+            logger.error("Error downloading file using signed URL: %s", e)
+            raise
+
+    def download_with_boto3(self, remote_file_path: str, local_file_path: str):
+        """Download a file using boto3 and credentials."""
+        try:
+            self.s3_client.download_file(self.dataset_config.s3.bucket, remote_file_path, local_file_path)
+            logger.info("File downloaded successfully to %s using S3 client.", local_file_path)
+        except Exception as e:
+            logger.error("Error downloading file from S3: %s", e)
+            raise
+
+    @staticmethod
+    def is_file_path_url(file_path: str) -> bool:
+        """Check if the file path is a URL."""
+        return file_path.startswith("http")
+
+    def download_file(self, remote_file_path: str, local_file_path: str):
+        """Download a file using the appropriate method."""
+        self.ensure_directory_exists(local_file_path)
+        if self.is_file_path_url(remote_file_path):
+            logger.info("Using signed URL to download the file %s...", remote_file_path)
+            self.download_with_signed_url(remote_file_path, local_file_path, timeout=120)
+        else:
+            logger.info("Using S3 credentials to download the file %s...", remote_file_path)
+            self.download_with_boto3(remote_file_path, local_file_path)
+
+    def download_dataset(self):
+        """Download datasets defined in the evaluation configuration."""
+        if self.dataset_config.remote_file_path:
+            logger.info("Downloading remote dataset %s")
+            self.download_file(remote_file_path=self.dataset_config.remote_file_path,
+                               local_file_path=self.dataset_config.file_path)

nat/eval/dataset_handler/dataset_filter.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pandas as pd
+
+from nat.data_models.dataset_handler import EvalFilterConfig
+
+
+class DatasetFilter:
+    """
+    Apply allowlist and denylist filters to the DataFrame based on specified column filters.
+        - If a allowlist is provided, only keep rows matching the filter values.
+        - If a denylist is provided, remove rows matching the filter values.
+        - If the filter column does not exist in the DataFrame, the filtering is skipped for that column.
+
+    This is a utility class that is dataset agnostic and can be used to filter any DataFrame based on the provided
+    filter configuration.
+    """
+
+    def __init__(self, filter_config: EvalFilterConfig):
+
+        self.filter_config = filter_config
+
+    def apply_filters(self, df) -> pd.DataFrame:
+
+        filtered_df = df.copy()
+
+        # Apply allowlist (only keep specified rows)
+        if self.filter_config.allowlist:
+            for column, values in self.filter_config.allowlist.field.items():
+                if column in filtered_df.columns:
+                    filtered_df = filtered_df[filtered_df[column].isin(values)]
+
+        # Apply denylist (remove specified rows)
+        if self.filter_config.denylist:
+            for column, values in self.filter_config.denylist.field.items():
+                if column in filtered_df.columns:
+                    filtered_df = filtered_df[~filtered_df[column].isin(values)]
+
+        return filtered_df