nvidia-nat 1.4.0a20251112__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/experimental/test_time_compute/selection/llm_based_plan_selector.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/experimental/test_time_compute/selection/llm_judge_selection.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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
+import re
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.experimental.test_time_compute.models.selection_config import LLMJudgeSelectionConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
+
+logger = logging.getLogger(__name__)
+
+
+class LLMJudgeSelection(StrategyBase):
+    """
+    A selection strategy that uses a configured Judge LLM to select the best response.
+    """
+
+    def __init__(self, config: LLMJudgeSelectionConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self.judge_llm_bound = None
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Builds the Judge LLM configured in the strategy.
+        """
+        logger.debug("Building components for LLMJudgeSelection")
+        self.judge_llm_bound = await builder.get_llm(self.config.judge_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> list[PipelineTypeEnum]:
+        return [PipelineTypeEnum.CUSTOM, PipelineTypeEnum.PLANNING, PipelineTypeEnum.AGENT_EXECUTION]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SELECTION
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        Select the best item using the configured Judge LLM.
+        """
+        if not self.judge_llm_bound:
+            raise ValueError("Judge LLM not bound. Ensure `build_components` has been called.")
+
+        if not items:
+            logger.warning("No items provided for selection.")
+            return []
+
+        try:
+            from langchain_core.prompts import PromptTemplate
+            from pydantic import BaseModel
+        except ImportError as exc:
+            raise ImportError("langchain-core is not installed.") from exc
+
+        # Format the results for the prompt
+        results_str = ""
+        for idx, item in enumerate(items):
+            item_output = (str(item.output.model_dump()) if isinstance(item.output, BaseModel) else str(item.output))
+            results_str += f"{idx + 1}. {remove_r1_think_tags(item_output)}\n\n"
+
+        prompt_template = PromptTemplate(
+            template=self.config.selection_template,
+            input_variables=["original_prompt", "results"],
+            validate_template=True,
+        )
+
+        # Use input from first item if original_prompt is missing
+        query = original_prompt if original_prompt else (items[0].input or "Unknown Query")
+
+        prompt = (await prompt_template.ainvoke(input={"original_prompt": query, "results": results_str})).to_string()
+
+        logger.info("Asking Judge LLM to select the best response.")
+        judge_response = await self.judge_llm_bound.ainvoke(prompt)
+        judge_content = remove_r1_think_tags(
+            judge_response.content if hasattr(judge_response, 'content') else str(judge_response))
+
+        # Parse selection
+        # Expected format: 'SELECTED ITEM: <number>'
+        match = re.search(r'SELECTED ITEM:\s*(\d+)', judge_content, re.IGNORECASE)
+        if match:
+            try:
+                index = int(match.group(1)) - 1
+                if 0 <= index < len(items):
+                    logger.info("Judge selected item %d", index + 1)
+                    selected_item = items[index]
+                    # Optionally attach judge's reasoning to metadata
+                    if selected_item.metadata is None:
+                        selected_item.metadata = {}
+                    selected_item.metadata["judge_reasoning"] = judge_content
+                    return [selected_item]
+                else:
+                    logger.warning("Judge selected index %d which is out of range.", index + 1)
+            except ValueError:
+                logger.warning("Failed to parse integer from judge selection.")
+
+        logger.warning("Could not parse valid selection from judge response. "
+                       "Returning first item as fallback.")
+        # Fallback to first item
+        return [items[0]]
+
+
+@register_ttc_strategy(config_type=LLMJudgeSelectionConfig)
+async def register_llm_judge_selection(config: LLMJudgeSelectionConfig, builder: Builder):
+    strategy = LLMJudgeSelection(config)
+    await strategy.build_components(builder)
+    yield strategy

nat/experimental/test_time_compute/selection/threshold_selector.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/finetuning/__init__.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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 nat.finetuning.interfaces.finetuning_runner import Trainer
+from nat.finetuning.interfaces.trainer_adapter import TrainerAdapter
+from nat.finetuning.interfaces.trajectory_builder import TrajectoryBuilder
+
+__all__ = [
+    "Trainer",
+    "TrajectoryBuilder",
+    "TrainerAdapter",
+]

nat/finetuning/finetuning_runtime.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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.
+"""Finetuning runtime for NAT that orchestrates the training process."""
+
+import asyncio
+import logging
+
+from nat.data_models.finetuning import FinetuneRunConfig
+from nat.data_models.finetuning import TrainingStatusEnum
+from nat.finetuning.interfaces.finetuning_runner import Trainer
+
+logger = logging.getLogger(__name__)
+
+
+async def run_finetuning(runner: Trainer) -> None:
+    """
+    Run finetuning based on the provided configuration.
+
+    Args:
+        runner: An instance of the Trainer to run finetuning with
+    """
+    try:
+        # Initialize the runner
+        logger.info("Initializing finetuning runner...")
+
+        # Get number of epochs from config
+        num_epochs = runner.run_config.num_epochs
+
+        # Run training for specified epochs
+        logger.info("Starting training for %d epochs...", num_epochs)
+        job_statuses = await runner.run(num_epochs)
+
+        # Log final status
+        for status in job_statuses:
+            logger.info("Job %s completed with status: %s", status.run_id, status.status)
+            if status.message:
+                logger.info("  Message: %s", status.message)
+
+        # Get and log final metrics
+        if job_statuses:
+            final_run_id = job_statuses[-1].run_id
+            try:
+                metrics = await runner.get_metrics(final_run_id)
+                logger.info("Final metrics: %s", metrics)
+            except (ValueError, RuntimeError) as e:
+                logger.warning("Failed to retrieve metrics: %s", e)
+
+        # Log appropriate message based on job statuses
+        if not job_statuses:
+            logger.warning("Finetuning completed with no jobs executed.")
+        else:
+            failed_jobs = sum(1 for s in job_statuses if s.status == TrainingStatusEnum.FAILED)
+            canceled_jobs = sum(1 for s in job_statuses if s.status == TrainingStatusEnum.CANCELED)
+            completed_jobs = sum(1 for s in job_statuses if s.status == TrainingStatusEnum.COMPLETED)
+
+            if failed_jobs:
+                logger.error("Finetuning completed with %d failed job(s) out of %d total.",
+                             failed_jobs,
+                             len(job_statuses))
+            elif canceled_jobs:
+                logger.warning("Finetuning was canceled. %d job(s) were canceled out of %d total.",
+                               canceled_jobs,
+                               len(job_statuses))
+            elif completed_jobs == len(job_statuses):
+                logger.info("Finetuning completed successfully!")
+            else:
+                # Some jobs may still be pending or running (unexpected state)
+                logger.warning("Finetuning finished with %d completed, %d pending/running job(s).",
+                               completed_jobs,
+                               len(job_statuses) - completed_jobs)
+
+    except Exception as e:
+        logger.error("Finetuning failed: %s", e)
+        raise
+    finally:
+        # Always cleanup resources
+        logger.info("Cleaning up finetuning resources...")
+        await runner.cleanup()
+        logger.info("Cleanup completed")
+
+
+async def finetuning_main(run_config: FinetuneRunConfig) -> None:
+    """
+    Main entry point for finetuning runtime.
+
+    Args:
+        run_config: FinetuneRunConfig object containing finetuning settings
+    """
+
+    from nat.builder.workflow_builder import WorkflowBuilder
+    from nat.runtime.loader import load_config
+
+    config = load_config(config_file=run_config.config_file)
+    finetuning_config = config.finetuning
+    finetuning_config.run_configuration = run_config
+
+    if not config.finetuning.enabled:
+        raise ValueError("Finetuning is not enabled in the provided configuration.")
+
+    async with WorkflowBuilder.from_config(config=config) as builder:
+        # Get trajectory builder and trainer adapter from builder
+        logger.info("Initializing finetuning components...")
+        trajectory_builder_name = finetuning_config.trajectory_builder
+        trainer_adapter_name = finetuning_config.trainer_adapter
+        trajectory_builder = await builder.get_trajectory_builder(trajectory_builder_name)
+        trainer_adapter = await builder.get_trainer_adapter(trainer_adapter_name)
+        logger.info("Finetuning components initialized.")
+
+        # Initialize trainer
+        trainer_name = finetuning_config.trainer
+        trainer = await builder.get_trainer(trainer_name,
+                                            trajectory_builder=trajectory_builder,
+                                            trainer_adapter=trainer_adapter)
+
+        await trainer.initialize(run_config=finetuning_config)
+
+        logger.info("Initialized trainer: %s", trainer_name)
+
+    # Run finetuning
+    await run_finetuning(trainer)
+
+
+def run_finetuning_sync(run_config: FinetuneRunConfig) -> None:
+    """
+    Synchronous wrapper for running finetuning.
+
+    Args:
+        run_config: FinetuneRunConfig object containing finetuning settings
+    """
+    asyncio.run(finetuning_main(run_config))

nat/finetuning/interfaces/__init__.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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 nat.finetuning.interfaces.finetuning_runner import Trainer
+from nat.finetuning.interfaces.trainer_adapter import TrainerAdapter
+from nat.finetuning.interfaces.trajectory_builder import TrajectoryBuilder
+
+__all__ = [
+    "Trainer",
+    "TrajectoryBuilder",
+    "TrainerAdapter",
+]

nat/finetuning/interfaces/finetuning_runner.py

@@ -0,0 +1,261 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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 abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+from nat.data_models.finetuning import FinetuneConfig
+from nat.data_models.finetuning import FinetuneRunConfig
+from nat.data_models.finetuning import TrainerConfig
+from nat.data_models.finetuning import TrainingJobRef
+from nat.data_models.finetuning import TrainingJobStatus
+from nat.data_models.finetuning import TrajectoryCollection
+from nat.eval.config import EvaluationRunOutput
+from nat.finetuning.interfaces.trainer_adapter import TrainerAdapter
+from nat.finetuning.interfaces.trajectory_builder import TrajectoryBuilder
+
+logger = logging.getLogger(__name__)
+
+
+class Trainer(ABC):
+    """
+    Abstract interface for running finetuning workflows.
+
+    The Trainer orchestrates the entire finetuning process by:
+    1. Running evaluations to generate trajectories via TrajectoryBuilder
+    2. Submitting trajectories for training via TrainerAdapter
+    3. Managing multiple epochs of training
+    """
+
+    def __init__(self, trainer_config: TrainerConfig, **kwargs) -> None:
+        """
+        Initialize the Trainer.
+
+        Args:
+            trainer_config: Configuration for the trainer backend
+            run_config: Configuration for the training run
+            backend: Backend identifier
+            curriculum_config: Optional curriculum learning configuration
+        """
+        self.trainer_config = trainer_config
+        self.run_config: FinetuneConfig = None
+        self.curriculum_config = None
+        self.trajectory_builder: TrajectoryBuilder = None
+        self.trainer_adapter: TrainerAdapter = None
+
+        # Curriculum learning state
+        self._curriculum_state = None
+
+    async def bind_components(self, trajectory_builder: TrajectoryBuilder, trainer_adapter: TrainerAdapter) -> None:
+        """
+        Bind the TrajectoryBuilder and TrainerAdapter components.
+
+        Args:
+            trajectory_builder: Instance of TrajectoryBuilder
+            trainer_adapter: Instance of TrainerAdapter
+        """
+        self.trajectory_builder = trajectory_builder
+        self.trainer_adapter = trainer_adapter
+
+    async def initialize(self, run_config: FinetuneConfig) -> None:
+        """
+        Initialize the runner and its components.
+
+        This should:
+        - Initialize the TrajectoryBuilder
+        - Initialize the TrainerAdapter
+        - Verify connectivity to backend services
+        """
+
+        self.run_config = run_config
+        self.curriculum_config = self.run_config.curriculum_learning
+        self._curriculum_state = {
+            "current_percentile": self.curriculum_config.initial_percentile,
+            "last_expansion_epoch": -1,
+            "total_groups": 0,
+            "included_groups": set()
+        }
+        self.trainer_config.reward = self.run_config.reward_function
+
+        await self.trajectory_builder.initialize(run_config)
+        await self.trainer_adapter.initialize(run_config)
+
+    @abstractmethod
+    async def run_epoch(self, epoch: int, run_id: str) -> TrainingJobRef:
+        """
+        Run a single epoch of training.
+
+        Args:
+            epoch: The current epoch number (0-indexed)
+            run_id: Unique identifier for this training run
+
+        Returns:
+            TrainingJobRef: Reference to the submitted training job
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def run(self, num_epochs: int) -> list[TrainingJobStatus]:
+        """
+        Run the complete finetuning workflow for the specified number of epochs.
+
+        Args:
+            num_epochs: Number of epochs to train
+
+        Returns:
+            list[TrainingJobStatus]: Status of all training jobs
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get_metrics(self, run_id: str) -> dict[str, Any]:
+        """
+        Get training metrics for a specific run.
+
+        Args:
+            run_id: The run identifier
+
+        Returns:
+            dict: Metrics from the training run
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def cleanup(self) -> None:
+        """
+        Clean up any resources used by the runner.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def log_progress(self, epoch: int, metrics: dict[str, Any], output_dir: str | None = None) -> None:
+        """
+        Log training progress for monitoring.
+
+        Args:
+            epoch: Current epoch number
+            metrics: Dictionary of metrics to log
+            output_dir: Optional output directory override
+        """
+        raise NotImplementedError
+
+    async def run_validation_evaluation(self, epoch: int, run_id: str) -> dict[str, Any]:
+        """
+        Run evaluation on validation dataset to collect rewards.
+
+        This method creates a temporary TrainerRunConfig with the validation
+        dataset and runs evaluation to collect rewards without training.
+
+        Args:
+            epoch: Current epoch number
+            run_id: Unique identifier for this training run
+            validation_dataset: Path to the validation dataset
+
+        Returns:
+            dict: Validation metrics including average reward
+        """
+        logger.info("Running validation evaluation for epoch %d", epoch + 1)
+
+        config = self.run_config.run_configuration.validation_config_file if (
+            self.run_config.run_configuration.validation_config_file) else self.run_config.run_configuration.config_file
+
+        # Create a temporary run config with validation dataset
+        validation_run_config = FinetuneRunConfig(config_file=config,
+                                                  dataset=self.run_config.run_configuration.validation_dataset,
+                                                  result_json_path=self.run_config.run_configuration.result_json_path,
+                                                  endpoint=self.run_config.run_configuration.endpoint,
+                                                  endpoint_timeout=self.run_config.run_configuration.endpoint_timeout,
+                                                  override=self.run_config.run_configuration.override)
+
+        # Create a temporary trajectory builder for validation
+        validation_builder = self.trajectory_builder
+        original_run_config = validation_builder.run_config.run_configuration
+
+        try:
+
+            validation_builder.run_config.run_configuration = validation_run_config
+
+            # Run evaluation
+            eval_output = await validation_builder.run_eval()
+
+            # Calculate validation metrics from eval output
+            validation_metrics = self._calculate_validation_metrics(eval_output)
+            validation_metrics["epoch"] = epoch
+            validation_metrics["dataset_type"] = "validation"
+
+            logger.info("Validation metrics for epoch %d: %s", epoch, validation_metrics)
+            return validation_metrics
+
+        except Exception as e:
+            logger.error("Error during validation evaluation: %s", e)
+            return {"epoch": epoch, "dataset_type": "validation", "error": str(e), "avg_reward": 0.0, "num_examples": 0}
+        finally:
+            # Restore original run config
+            validation_builder.run_config.run_configuration = original_run_config
+
+    def _calculate_validation_metrics(self, eval_output: EvaluationRunOutput) -> dict[str, Any]:
+        """
+        Calculate validation metrics from evaluation output.
+
+        Args:
+            eval_output: Output from evaluation run
+
+        Returns:
+            dict: Calculated metrics
+        """
+        # Default implementation - subclasses can override for
+        # backend-specific metrics
+        metrics = {"avg_reward": 0.0, "min_reward": 0.0, "max_reward": 0.0, "num_examples": 0}
+
+        rewards = []
+        for metric_name, metric_value in eval_output.evaluation_results:
+            if metric_name == self.trainer_config.reward.name:
+                reward_results = metric_value.eval_output_items
+                for reward_item in reward_results:
+                    rewards.append(reward_item.score)
+
+        if rewards:
+            metrics["avg_reward"] = sum(rewards) / len(rewards)
+            metrics["min_reward"] = min(rewards)
+            metrics["max_reward"] = max(rewards)
+            metrics["num_examples"] = len(rewards)
+
+        return metrics
+
+    def apply_curriculum_learning(self, trajectory_collection: TrajectoryCollection,
+                                  epoch: int) -> TrajectoryCollection:
+        """
+        Apply curriculum learning to filter trajectory groups based on difficulty.
+        """
+        raise NotImplementedError("Curriculum learning not implemented for this backend.")
+
+    def get_curriculum_state(self) -> dict[str, Any]:
+        """
+        Get the current state of curriculum learning.
+
+        Returns:
+            dict: Current curriculum state including percentile and group statistics
+        """
+        # Convert set to list for JSON serialization
+        state = {
+            "current_percentile": self._curriculum_state["current_percentile"],
+            "last_expansion_epoch": self._curriculum_state["last_expansion_epoch"],
+            "total_groups": self._curriculum_state["total_groups"],
+            "included_groups": list(self._curriculum_state["included_groups"]),
+            "config": self.curriculum_config.model_dump() if self.curriculum_config else None
+        }
+        return state