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

nat/finetuning/interfaces/trainer_adapter.py

@@ -0,0 +1,103 @@
+# 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 abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+from nat.data_models.finetuning import FinetuneConfig
+from nat.data_models.finetuning import TrainerAdapterConfig
+from nat.data_models.finetuning import TrainingJobRef
+from nat.data_models.finetuning import TrainingJobStatus
+from nat.data_models.finetuning import TrajectoryCollection
+
+
+class TrainerAdapter(ABC):
+    """
+    Adapter to send Trajectories to remote training cluster for weights updates.
+    """
+
+    def __init__(self, adapter_config: TrainerAdapterConfig):
+        self.adapter_config = adapter_config
+        self.run_config: FinetuneConfig = None
+
+    async def initialize(self, run_config: FinetuneConfig) -> None:
+        """
+        Asynchronously initialize any resources needed for the trainer adapter.
+        """
+        self.run_config = run_config
+        self.adapter_config.reward = self.run_config.reward_function
+
+    @abstractmethod
+    async def is_healthy(self) -> bool:
+        """
+        Check the health of the remote training backend.
+
+        Returns:
+            bool: True if the backend is healthy, False otherwise.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def submit(self, trajectories: TrajectoryCollection) -> TrainingJobRef:
+        """
+        Submit trajectories to remote training backend.
+
+        Args:
+            trajectories (list[Trajectory]): The list of trajectories to submit.
+
+        Returns:
+            TrainingJobRef: Reference to the submitted training job.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def status(self, ref: TrainingJobRef) -> TrainingJobStatus:
+        """
+        Get the status of a submitted training job.
+
+        Args:
+            ref (TrainingJobRef): Reference to the training job.
+
+        Returns:
+            TrainingJobStatus: The current status of the training job.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def wait_until_complete(self, ref: TrainingJobRef, poll_interval: float = 10.0) -> TrainingJobStatus:
+        """
+        Wait until the training job is complete.
+
+        Args:
+            ref (TrainingJobRef): Reference to the training job.
+            poll_interval (float): Time in seconds between status checks.
+
+        Returns:
+            TrainingJobStatus: The final status of the training job.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def log_progress(self, ref: TrainingJobRef, metrics: dict[str, Any], output_dir: str | None = None) -> None:
+        """
+        Log training adapter progress.
+
+        Args:
+            ref: Training job reference
+            metrics: Dictionary of metrics to log
+            output_dir: Optional output directory override
+        """
+        raise NotImplementedError

nat/finetuning/interfaces/trajectory_builder.py

@@ -0,0 +1,115 @@
+# 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 abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+from nat.data_models.finetuning import FinetuneConfig
+from nat.data_models.finetuning import TrajectoryBuilderConfig
+from nat.data_models.finetuning import TrajectoryCollection
+from nat.eval.config import EvaluationRunOutput
+from nat.eval.evaluator.evaluator_model import EvalOutputItem
+from nat.utils.io.supress_logs import suppress_logs
+
+
+class TrajectoryBuilder(ABC):
+    """
+    Abstract interface for building trajectories from episode items.
+    """
+
+    def __init__(self, trajectory_builder_config: TrajectoryBuilderConfig):
+        self.trajectory_builder_config = trajectory_builder_config
+        self.run_config: FinetuneConfig = None
+
+    async def initialize(self, run_config: FinetuneConfig) -> None:
+        """
+        Asynchronously initialize any resources needed for the trajectory builder.
+        """
+        self.run_config = run_config
+        self.trajectory_builder_config.reward = self.run_config.reward_function
+
+    async def run_eval(self) -> EvaluationRunOutput:
+        """
+        Run NAT Evaluation to generate episode items for trajectory building.
+
+        Returns:
+            EvaluationRunOutput: The output of the evaluation run.
+        """
+
+        from nat.eval.evaluate import EvaluationRun
+        from nat.eval.evaluate import EvaluationRunConfig
+
+        eval_cfg = EvaluationRunConfig(config_file=self.run_config.run_configuration.config_file,
+                                       dataset=self.run_config.run_configuration.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)
+
+        async with suppress_logs(prefix="nat.eval"):
+            evaluation_output = await EvaluationRun(config=eval_cfg).run_and_evaluate()
+
+        return evaluation_output
+
+    @abstractmethod
+    async def start_run(self, run_id: str, meta: dict | None = None) -> None:
+        """
+        Initialize any resources needed for the trajectory builder.
+
+        Args:
+            run_id (str): The unique identifier for the training run.
+            meta (dict): Metadata associated with the training run.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def finalize(self, run_id: str, meta: dict | None = None) -> TrajectoryCollection:
+        """
+        Finalize the trajectory building process and return the constructed trajectories.
+
+        Args:
+            run_id (str): The unique identifier for the training run.
+            meta (dict): Metadata associated with the training run.
+
+        Returns:
+            list[Trajectory]: The list of constructed trajectories.
+        """
+        raise NotImplementedError
+
+    async def compute_reward(self, output_item: EvalOutputItem, meta: dict | None = None):
+        """
+        Compute reward for a given EvalOutputItem.
+
+        Args:
+            output_item (EvalOutputItem): The evaluation output item.
+            meta (dict): Metadata associated with the training run.
+
+        Returns:
+            float: The computed reward.
+        """
+        return float(output_item.score) if output_item.score is not None else 0.0
+
+    @abstractmethod
+    def log_progress(self, run_id: str, metrics: dict[str, Any], output_dir: str | None = None) -> None:
+        """
+        Log trajectory building progress.
+
+        Args:
+            run_id: The training run ID
+            metrics: Dictionary of metrics to log
+            output_dir: Optional output directory override
+        """
+        raise NotImplementedError

nat/finetuning/utils/__init__.py

@@ -0,0 +1,15 @@
+# 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.

nat/finetuning/utils/parsers/__init__.py

@@ -0,0 +1,15 @@
+# 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.

nat/finetuning/utils/parsers/adk_parser.py

@@ -0,0 +1,141 @@
+# 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 json
+import logging
+
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.finetuning.utils.parsers.common import extract_content
+from nat.finetuning.utils.parsers.common import parse_generic_message
+
+logger = logging.getLogger(__name__)
+
+# Re-export for backwards compatibility and internal use
+_extract_content = extract_content
+_parse_generic_message = parse_generic_message
+
+
+def parse_to_openai_message(message: IntermediateStep) -> dict:  # noqa: ARG001
+    """
+    Convert IntermediateStep to OpenAI-compatible message dictionary.
+
+    Args:
+        message: An IntermediateStep object representing a single message.
+        previous_message: Previous message for context (reserved for future).
+
+    Returns:
+        A dictionary formatted for OpenAI API consumption.
+    """
+    result = {}
+
+    # Handle different event types to determine role and extract content
+    if message.event_type == IntermediateStepType.LLM_END:
+        # Assistant message from ChatResponse
+        result = _parse_assistant_message(message)
+    elif message.event_type == IntermediateStepType.TOOL_END:
+        # Tool/Function response message
+        result = _parse_tool_message(message)
+    elif message.event_type == IntermediateStepType.LLM_START:
+        # Extract user/system messages from the input
+        result = _parse_input_message(message)
+    else:
+        # For other types, try to infer from the data
+        result = _parse_generic_message(message)
+
+    return result
+
+
+def _parse_input_message(message: IntermediateStep) -> dict | list[dict]:
+    """Parse user or system messages from LLM_START event."""
+
+    messages = message.data.payload
+
+    if len(messages) == 0:
+        return {"role": "user", "content": ""}
+    elif len(messages) == 1:
+        if not isinstance(messages[0], dict):
+            return {"role": "user", "content": str(messages[0])}
+
+        if not ("role" in messages[0] and "content" in messages[0]):
+            return {"role": "user", "content": json.dumps(messages[0])}
+
+        return messages[0]
+    else:
+        parsed_messages = []
+        for msg in messages:
+            if not isinstance(msg, dict):
+                parsed_messages.append({"role": "user", "content": str(msg)})
+            elif not ("role" in msg and "content" in msg):
+                parsed_messages.append({"role": "user", "content": json.dumps(msg)})
+            else:
+                parsed_messages.append(msg)
+        return parsed_messages
+
+
+def _parse_assistant_message(message: IntermediateStep) -> dict:
+    """Parse an assistant message from LLM_END event."""
+    result = {"role": "assistant"}
+
+    # Get the ChatResponse from payload if available
+    try:
+        if message.data and message.data.payload:
+            pass
+            payload = message.data.payload
+            payload_message = getattr(payload, 'message', None)
+
+            if "logprobs" in payload:
+                result["logprobs"] = payload["logprobs"]
+            else:
+                logger.warning("No logprobs found in LLM_END message payload.")
+
+            if "content" in payload_message and payload_message["content"] is not None:
+                result["content"] = _extract_content(payload_message["content"])
+            else:
+                result["content"] = ""
+
+            if "tool_calls" in payload_message and payload_message["tool_calls"] is not None:
+                result["tool_calls"] = payload_message["tool_calls"]
+
+        else:
+            logger.warning("No payload found in LLM_END message data.")
+            return {"role": "assistant", "content": ""}
+    except Exception as _:
+        logger.exception("Error parsing assistant message from LLM_END event.")
+        return {"role": "assistant", "content": ""}
+
+    return result
+
+
+def _parse_tool_message(message: IntermediateStep) -> dict:
+    """Parse a tool/function response message from TOOL_END event."""
+    result = {"role": "function"}
+
+    # Extract function output as content
+    if message.data:
+        if message.data.output:
+            result["content"] = _extract_content(message.data.output)
+        elif message.data.payload:
+            result["content"] = _extract_content(message.data.payload)
+        else:
+            result["content"] = ""
+    else:
+        result["content"] = ""
+
+    # Add function name if available
+    if message.name:
+        result["name"] = message.name
+
+    return result

nat/finetuning/utils/parsers/base_parser.py

@@ -0,0 +1,238 @@
+# 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 nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.data_models.intermediate_step import IntermediateStepState
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.finetuning.utils.parsers import adk_parser
+from nat.finetuning.utils.parsers import langchain_parser
+from nat.finetuning.utils.parsers import llama_index_parser
+
+logger = logging.getLogger(__name__)
+
+
+def parse_to_openai_messages(steps: list[IntermediateStep]) -> list[dict]:
+    """
+    Convert IntermediateStep objects to OpenAI-compatible messages.
+
+    Args:
+        steps: List of IntermediateStep objects representing the conversation.
+
+    Returns:
+        List of dictionaries formatted for OpenAI API consumption.
+
+    Raises:
+        ValueError: If unsupported type or invalid sequence.
+    """
+
+    messages = []
+
+    # Track the last event type to handle special cases
+    last_event_type = None
+    message_content_hashes = set()
+    for message in steps:
+        # Skip LLM_START events that come after TOOL_END events
+        # These represent the assistant processing tool results internally
+        if message.event_type not in [
+                IntermediateStepType.LLM_END, IntermediateStepType.LLM_START, IntermediateStepType.TOOL_END
+        ]:
+            continue
+
+        if (message.event_type == IntermediateStepType.LLM_START and last_event_type == IntermediateStepType.TOOL_END):
+            continue
+
+        # Skip streaming chunks
+        if message.event_state not in [IntermediateStepState.START, IntermediateStepState.END]:
+            continue
+
+        # Parse the message based on framework
+        if message.framework == LLMFrameworkEnum.LANGCHAIN:
+            parsed_msg = langchain_parser.parse_to_openai_message(message=message)
+        elif message.framework == LLMFrameworkEnum.LLAMA_INDEX:
+            parsed_msg = llama_index_parser.parse_to_openai_message(message=message)
+        elif message.framework == LLMFrameworkEnum.ADK:
+            parsed_msg = adk_parser.parse_to_openai_message(message=message)
+        else:
+            if message.framework is not None:
+                logger.warning(f"Unsupported framework: {message.framework} for message {message}")
+            continue
+
+        # Add the parsed message
+        if message.event_type == IntermediateStepType.LLM_START:
+            # LLM_START messages may contain multiple messages (e.g., tools called by the LLM)
+            # We deduplicate previously seen messages if sharing message history to the model
+            if isinstance(parsed_msg, list):
+                for msg in parsed_msg:
+                    content_hash = hash(msg["role"] + ": " + msg["content"])
+                    if content_hash not in message_content_hashes:
+                        messages.append(msg)
+                        message_content_hashes.add(content_hash)
+            else:
+                content_hash = hash(parsed_msg["role"] + ": " + parsed_msg["content"])
+                messages.append(parsed_msg)
+                message_content_hashes.add(content_hash)
+        else:
+            assert not isinstance(parsed_msg, list), "TOOL_END or LLM_END should not produce multiple messages"
+            content_hash = hash(parsed_msg["role"] + ": " + parsed_msg["content"])
+            message_content_hashes.add(content_hash)
+            messages.append(parsed_msg)
+
+        last_event_type = message.event_type
+
+    # Validate and fix the message sequence
+    try:
+        messages = _validate_message_sequence(messages)
+    except Exception as _:
+        logger.exception("Error validating message sequence.")
+        raise
+
+    return messages
+
+
+def _validate_message_sequence(messages: list[dict]) -> list[dict]:
+    """
+    Validate and fix the message sequence to follow OpenAI's expected format.
+
+    Rules:
+
+    - System messages can only appear at the beginning
+    - After system messages, must alternate between user/tool and assistant
+    - Cannot have consecutive user messages or consecutive assistant messages
+    - If first non-system messages are not user messages, they will be
+      concatenated into a single user message (with a warning)
+
+    Args:
+        messages: List of parsed OpenAI messages
+
+    Returns:
+        list[dict]: The validated (and potentially fixed) message list
+
+    Raises:
+        ValueError: If the message sequence is invalid.
+    """
+    if not messages:
+        return messages
+
+    # Check system messages are only at the beginning
+    found_non_system = False
+    for i, msg in enumerate(messages):
+        if msg.get("role") == "system":
+            if found_non_system:
+                raise ValueError(f"System message found at position {i} after "
+                                 "non-system messages. System messages must only "
+                                 "appear at the beginning.")
+        else:
+            found_non_system = True
+
+    # Find first non-system message
+    first_non_system_idx = 0
+    for i, msg in enumerate(messages):
+        if msg.get("role") != "system":
+            first_non_system_idx = i
+            break
+
+    # Fix non-user messages at the start of trajectory
+    # Collect all non-system messages before the first assistant message
+    if first_non_system_idx < len(messages):
+        # Find the first assistant message
+        first_assistant_idx = None
+        for i in range(first_non_system_idx, len(messages)):
+            if messages[i].get("role") == "assistant":
+                first_assistant_idx = i
+                break
+
+        # Check if we need to fix the start of the trajectory
+        if first_assistant_idx is not None:
+            messages_to_concatenate = []
+            for i in range(first_non_system_idx, first_assistant_idx):
+                msg = messages[i]
+                role = msg.get("role")
+                if role != "user":
+                    # This message should be concatenated
+                    messages_to_concatenate.append((i, msg))
+
+            if messages_to_concatenate:
+                # Collect all content from non-user messages at the start
+                content_parts = []
+                indices_to_remove = []
+
+                for i in range(first_non_system_idx, first_assistant_idx):
+                    msg = messages[i]
+                    role = msg.get("role")
+                    content = msg.get("content", "")
+
+                    if role not in ["user"]:
+                        # Non-user message that needs to be consolidated
+                        if content:
+                            content_parts.append(f"[{role.upper()}]: {content}")
+                        indices_to_remove.append(i)
+                    else:
+                        # User message - include its content
+                        if content:
+                            content_parts.append(content)
+                        indices_to_remove.append(i)
+
+                # Create a single user message with concatenated content
+                if content_parts:
+                    concatenated_content = "\n\n".join(content_parts)
+                    new_user_message = {"role": "user", "content": concatenated_content}
+
+                    # Log warning about the modification
+                    logger.warning(
+                        "Trajectory had %d non-user messages at the start "
+                        "before the first assistant message. "
+                        "Concatenated these into a single user message. "
+                        "Original roles: %s",
+                        len(messages_to_concatenate), [msg.get("role") for _, msg in messages_to_concatenate])
+
+                    # Remove the old messages and insert the new one
+                    # Remove in reverse order to maintain indices
+                    for idx in reversed(indices_to_remove):
+                        messages.pop(idx)
+
+                    # Insert the new user message
+                    messages.insert(first_non_system_idx, new_user_message)
+
+    # Recalculate first_non_system_idx after potential modifications
+    first_non_system_idx = 0
+    for i, msg in enumerate(messages):
+        if msg.get("role") != "system":
+            first_non_system_idx = i
+            break
+
+    # Validate alternating pattern after system messages
+    if first_non_system_idx < len(messages):
+        prev_role = None
+        for i in range(first_non_system_idx, len(messages)):
+            role = messages[i].get("role")
+
+            if prev_role:
+                # Check for invalid consecutive roles
+                if role == "user" and prev_role == "user":
+                    raise ValueError(f"Consecutive user messages at positions {i-1} "
+                                     f"and {i}. User messages must be followed by "
+                                     "assistant messages.")
+                elif role == "assistant" and prev_role == "assistant":
+                    raise ValueError(f"Consecutive assistant messages at positions "
+                                     f"{i-1} and {i}. Assistant messages must be "
+                                     "followed by user or tool messages.")
+
+            prev_role = role
+
+    return messages