nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/data_models/optimizer.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2021-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 pydantic import Field
+
+
+class OptimizerMetric(BaseModel):
+    """
+    Parameters used by the workflow optimizer to define a metric to optimize.
+    """
+    evaluator_name: str = Field(description="Name of the metric to optimize.")
+    direction: str = Field(description="Direction of the optimization. Can be 'maximize' or 'minimize'.")
+    weight: float = Field(description="Weight of the metric in the optimization process.", default=1.0)
+
+
+class NumericOptimizationConfig(BaseModel):
+    """
+    Configuration for numeric/enum optimization (Optuna).
+    """
+    enabled: bool = Field(default=True, description="Enable numeric optimization")
+    n_trials: int = Field(description="Number of trials for numeric optimization.", default=20)
+
+
+class PromptGAOptimizationConfig(BaseModel):
+    """
+    Configuration for prompt optimization using a Genetic Algorithm.
+    """
+    enabled: bool = Field(default=False, description="Enable GA-based prompt optimization")
+
+    # Prompt optimization function hooks
+    prompt_population_init_function: str | None = Field(
+        default=None,
+        description="Optional function name to initialize/mutate candidate prompts.",
+    )
+    prompt_recombination_function: str | None = Field(
+        default=None,
+        description="Optional function name to recombine two parent prompts into a child.",
+    )
+
+    # Genetic algorithm configuration
+    ga_population_size: int = Field(
+        description="Population size for genetic algorithm prompt optimization.",
+        default=24,
+    )
+    ga_generations: int = Field(
+        description="Number of generations to evolve in GA prompt optimization.",
+        default=15,
+    )
+    ga_offspring_size: int | None = Field(
+        description="Number of offspring to produce per generation. Defaults to population_size - elitism.",
+        default=None,
+    )
+    ga_crossover_rate: float = Field(
+        description="Probability of applying crossover during reproduction.",
+        default=0.8,
+        ge=0.0,
+        le=1.0,
+    )
+    ga_mutation_rate: float = Field(
+        description="Probability of mutating a child after crossover.",
+        default=0.3,
+        ge=0.0,
+        le=1.0,
+    )
+    ga_elitism: int = Field(
+        description="Number of top individuals carried over unchanged each generation.",
+        default=2,
+    )
+    ga_selection_method: str = Field(
+        description="Parent selection strategy: 'tournament' or 'roulette'.",
+        default="tournament",
+    )
+    ga_tournament_size: int = Field(
+        description="Tournament size when using tournament selection.",
+        default=3,
+    )
+    ga_parallel_evaluations: int = Field(
+        description="Max number of individuals to evaluate concurrently per generation.",
+        default=8,
+    )
+    ga_diversity_lambda: float = Field(
+        description="Strength of diversity penalty (0 disables). Penalizes identical/near-identical prompts.",
+        default=0.0,
+        ge=0.0,
+    )
+
+
+class OptimizerConfig(BaseModel):
+    """
+    Parameters used by the workflow optimizer.
+    """
+    output_path: Path | None = Field(
+        default=None,
+        description="Path to the output directory where the results will be saved.",
+    )
+
+    eval_metrics: dict[str, OptimizerMetric] | None = Field(
+        description="List of evaluation metrics to optimize.",
+        default=None,
+    )
+
+    reps_per_param_set: int = Field(
+        default=3,
+        description="Number of repetitions per parameter set for the optimization.",
+    )
+
+    target: float | None = Field(
+        description=(
+            "Target value for the optimization. If set, the optimization will stop when this value is reached."),
+        default=None,
+    )
+
+    multi_objective_combination_mode: str = Field(
+        description="Method to combine multiple objectives into a single score.",
+        default="harmonic",
+    )
+
+    # Nested configs
+    numeric: NumericOptimizationConfig = NumericOptimizationConfig()
+    prompt: PromptGAOptimizationConfig = PromptGAOptimizationConfig()
+
+
+class OptimizerRunConfig(BaseModel):
+    """
+    Parameters used for an Optimizer R=run
+    """
+    # Eval parameters
+
+    config_file: Path | BaseModel  # allow for instantiated configs to be passed in
+    dataset: str | Path | None  # dataset file path can be specified in the config file
+    result_json_path: str = "$"
+    endpoint: str | None = None  # only used when running the workflow remotely
+    endpoint_timeout: int = 300
+    override: tuple[tuple[str, str], ...] = ()

nat/data_models/swe_bench_model.py

@@ -39,7 +39,7 @@ class SWEBenchInput(BaseModel):
 
     # Handle improperly formatted JSON strings for list fields
     @field_validator("FAIL_TO_PASS", "PASS_TO_PASS", mode="before")
-    def parse_list_fields(cls, value):  # pylint: disable=no-self-argument
+    def parse_list_fields(cls, value):
         if isinstance(value, str):
             # Attempt to parse the string as a list
             return json.loads(value)

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, 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:
+                    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/embedder/nim_embedder.py

@@ -50,7 +50,7 @@ class NIMEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="nim"):
                                        description=("The truncation strategy if the input on the "
                                                     "server side if it's too large."))
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
 
 @register_embedder_provider(config_type=NIMEmbedderModelConfig)

nat/embedder/openai_embedder.py

@@ -27,7 +27,7 @@ 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=())
+    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.")

nat/embedder/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

nat/eval/config.py

@@ -27,7 +27,7 @@ class EvaluationRunConfig(BaseModel):
     """
     Parameters used for a single evaluation run.
     """
-    config_file: Path
+    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
@@ -44,6 +44,8 @@ class EvaluationRunConfig(BaseModel):
     # 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):

nat/eval/dataset_handler/dataset_handler.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import json
 import math
 from pathlib import Path
@@ -41,7 +42,8 @@ class DatasetHandler:
                  reps: int,
                  concurrency: int,
                  num_passes: int = 1,
-                 adjust_dataset_size: bool = False):
+                 adjust_dataset_size: bool = False,
+                 custom_pre_eval_process_function: str | None = None):
         from nat.eval.intermediate_step_adapter import IntermediateStepAdapter
 
         self.dataset_config = dataset_config
@@ -53,6 +55,9 @@ class DatasetHandler:
         self.num_passes = num_passes
         self.adjust_dataset_size = adjust_dataset_size
 
+        # Custom pre-evaluation process function
+        self.custom_pre_eval_process_function = custom_pre_eval_process_function
+
         # Helpers
         self.intermediate_step_adapter = IntermediateStepAdapter()
 
@@ -146,13 +151,12 @@ class DatasetHandler:
             # When num_passes is specified, always use concurrency * num_passes
             # This respects the user's intent for exact number of passes
             target_size = self.concurrency * self.num_passes
+        # When num_passes = 0, use the largest multiple of concurrency <= original_size
+        # If original_size < concurrency, we need at least concurrency rows
+        elif original_size >= self.concurrency:
+            target_size = (original_size // self.concurrency) * self.concurrency
         else:
-            # When num_passes = 0, use the largest multiple of concurrency <= original_size
-            # If original_size < concurrency, we need at least concurrency rows
-            if original_size >= self.concurrency:
-                target_size = (original_size // self.concurrency) * self.concurrency
-            else:
-                target_size = self.concurrency
+            target_size = self.concurrency
 
         if target_size == 0:
             raise ValueError("Input dataset too small for even one batch at given concurrency.")
@@ -331,6 +335,66 @@ class DatasetHandler:
         filtered_steps = self.intermediate_step_adapter.filter_intermediate_steps(intermediate_steps, event_filter)
         return self.intermediate_step_adapter.serialize_intermediate_steps(filtered_steps)
 
+    def pre_eval_process_eval_input(self, eval_input: EvalInput) -> EvalInput:
+        """
+        Pre-evaluation process the eval input using custom function if provided.
+
+        The custom pre-evaluation process function should have the signature:
+        def custom_pre_eval_process(item: EvalInputItem) -> EvalInputItem
+
+        The framework will iterate through all items and call this function on each one.
+
+        Args:
+            eval_input: The EvalInput object to pre-evaluation process
+
+        Returns:
+            The pre-evaluation processed EvalInput object
+        """
+        if self.custom_pre_eval_process_function:
+            try:
+                custom_function = self._load_custom_pre_eval_process_function()
+                processed_items = []
+
+                for item in eval_input.eval_input_items:
+                    processed_item = custom_function(item)
+                    if not isinstance(processed_item, EvalInputItem):
+                        raise TypeError(f"Custom pre-evaluation '{self.custom_pre_eval_process_function}' must return "
+                                        f"EvalInputItem, got {type(processed_item)}")
+                    processed_items.append(processed_item)
+
+                return EvalInput(eval_input_items=processed_items)
+            except Exception as e:
+                raise RuntimeError(f"Error calling custom pre-evaluation process function "
+                                   f"'{self.custom_pre_eval_process_function}': {e}") from e
+
+        return eval_input
+
+    def _load_custom_pre_eval_process_function(self):
+        """
+        Import and return the custom pre-evaluation process function using standard Python import path.
+
+        The function should process individual EvalInputItem objects.
+        """
+        # Split the function path to get module and function name
+        if "." not in self.custom_pre_eval_process_function:
+            raise ValueError(f"Invalid custom_pre_eval_process_function '{self.custom_pre_eval_process_function}'. "
+                             "Expected format: '<module_path>.<function_name>'")
+        module_path, function_name = self.custom_pre_eval_process_function.rsplit(".", 1)
+
+        # Import the module
+        module = importlib.import_module(module_path)
+
+        # Get the function from the module
+        if not hasattr(module, function_name):
+            raise AttributeError(f"Function '{function_name}' not found in module '{module_path}'")
+
+        custom_function = getattr(module, function_name)
+
+        if not callable(custom_function):
+            raise ValueError(f"'{self.custom_pre_eval_process_function}' is not callable")
+
+        return custom_function
+
     def publish_eval_input(self,
                            eval_input,
                            workflow_output_step_filter: list[IntermediateStepType] | None = None) -> str: