From 9f651d4ad018a57c0f13a9bb3d752af151ed0b54 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Fri, 5 Sep 2025 14:15:49 -0700 Subject: [PATCH 1/5] Dataset viewer simple aggregations: - Display aggregate metrics - Aggregate generic keys using multineedle - Display other dynamic aggregations - Count string totals and unique values - Remove TrainDataProcessor dependency, add test - Remove dupe file read, fix arg types hints Signed-off-by: Frankie Siino --- nemo_gym/dataset_viewer.py | 121 ++++++++--- nemo_gym/train_data_utils.py | 289 ++++++++++++++++---------- tests/nemo_gym/test_dataset_viewer.py | 115 ++++++++++ 3 files changed, 377 insertions(+), 148 deletions(-) create mode 100644 tests/nemo_gym/test_dataset_viewer.py diff --git a/nemo_gym/dataset_viewer.py b/nemo_gym/dataset_viewer.py index 4ce8124028..93f229a127 100644 --- a/nemo_gym/dataset_viewer.py +++ b/nemo_gym/dataset_viewer.py @@ -1,33 +1,30 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# 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 typing import List, Dict, Any + import json -from typing import List -from gradio import Blocks, Chatbot, ChatMessage, Dropdown -from gradio.components.chatbot import MetadataDict +from tqdm.auto import tqdm + +from pydantic import BaseModel, ConfigDict + from openai.types.responses.response_input_param import ( - EasyInputMessageParam, - FunctionCallOutput, ResponseFunctionToolCallParam, - ResponseInputItemParam, + FunctionCallOutput, ResponseReasoningItemParam, + EasyInputMessageParam, + ResponseInputItemParam, ) -from pydantic import BaseModel, ConfigDict -from tqdm.auto import tqdm -from nemo_gym.base_resources_server import BaseVerifyResponse +from gradio import Chatbot, Blocks, ChatMessage, Dropdown, JSON +from gradio.components.chatbot import MetadataDict + from nemo_gym.server_utils import get_global_config_dict +from nemo_gym.base_resources_server import BaseVerifyResponse + +from nemo_gym.train_data_utils import ( + AvgMinMax, + compute_sample_metrics, + DatasetMetrics, +) class DatasetViewerVerifyResponse(BaseVerifyResponse): @@ -79,7 +76,9 @@ def format_reasoning(m: ResponseReasoningItemParam) -> List[ChatMessage]: def format_message(m: EasyInputMessageParam) -> List[ChatMessage]: - content = m["content"] if isinstance(m["content"], list) else [{"text": m["content"]}] + content = ( + m["content"] if isinstance(m["content"], list) else [{"text": m["content"]}] + ) match m["role"]: case "user": return [ @@ -131,7 +130,6 @@ def convert_single_message(m: ResponseInputItemParam) -> List[ChatMessage]: def rollout_to_messages(create_params: dict, response: dict) -> List[ChatMessage]: messages = [] - sampling_params = create_params.copy() sampling_params.pop("input") sampling_params.pop("tools", None) @@ -169,7 +167,9 @@ def rollout_to_messages(create_params: dict, response: dict) -> List[ChatMessage step += 1 for message in convert_single_message(m): - message.metadata["title"] = f"Turn {turn} Step {step} - {message.metadata['title']}" + message.metadata["title"] = ( + f"Turn {turn} Step {step} - {message.metadata['title']}" + ) messages.append(message) return messages @@ -202,16 +202,66 @@ class JsonlDatasetViewerConfig(BaseModel): jsonl_fpath: str +def aggregate_other_metrics(data: List[DatasetViewerVerifyResponse]) -> Dict[str, Any]: + metric_values = {} + string_values = {} + for d in data: + d = d.model_dump() if hasattr(d, "model_dump") else d + for k, v in d.items(): + if k in ("responses_create_params", "response"): + continue + if isinstance(v, bool): + v = int(v) + if isinstance(v, (int, float)): + metric_values.setdefault(k, []).append(v) + # get unique count for strings + elif isinstance(v, str): + string_values.setdefault(k, []).append(v) + + result = {} + for k, v in metric_values.items(): + if v: + obj = AvgMinMax( + total=len(v), + average=sum(v) / len(v), + min=min(v), + max=max(v), + ) + result[k] = obj.model_dump(by_alias=True) + + for k, v in string_values.items(): + result[k] = {"unique_count": len(set(v)), "total_count": len(v)} + + return result + + +def get_aggregate_metrics( + data: List[DatasetViewerVerifyResponse], raw_lines: List[str] +) -> Dict[str, Any]: + dataset_metrics = DatasetMetrics() + for line in raw_lines: + metrics, is_offending = compute_sample_metrics(line) + if not is_offending: + dataset_metrics.add(metrics) + + aggregate_metrics = dataset_metrics.aggregate() + aggregate_metrics_dict = aggregate_metrics.model_dump(by_alias=True) + aggregate_metrics_dict.update(**aggregate_other_metrics(data)) + return aggregate_metrics_dict + + def build_jsonl_dataset_viewer(config: JsonlDatasetViewerConfig) -> Blocks: + data = [] + raw_lines = [] with open(config.jsonl_fpath) as f: - data = list( - tqdm( - map(DatasetViewerVerifyResponse.model_validate_json, f), - desc="Loading data", - ) - ) + for line in tqdm(f, desc="Loading data"): + raw_lines.append(line) + data.append(DatasetViewerVerifyResponse.model_validate_json(line)) - choices = [(f"Sample {i + 1} - Responses ID {d.response.id}", i) for i, d in enumerate(data)] + choices = [ + (f"Sample {i + 1} - Responses ID {d.response.id}", i) + for i, d in enumerate(data) + ] def select_item(value: int): d = data[value] @@ -225,6 +275,9 @@ def select_item(value: int): } """ with Blocks(analytics_enabled=False, css=CSS) as demo: + aggregate_dicts = get_aggregate_metrics(data, raw_lines) + JSON(value=aggregate_dicts, label="Aggregate Metrics", open=False) + item_dropdown = Dropdown(choices=choices, value=0, label="Samples") chatbot = Chatbot( value=select_item(0), @@ -233,7 +286,9 @@ def select_item(value: int): layout="panel", label="Rollout", ) - item_dropdown.select(fn=select_item, inputs=item_dropdown, outputs=chatbot, show_api=False) + item_dropdown.select( + fn=select_item, inputs=item_dropdown, outputs=chatbot, show_api=False + ) return demo diff --git a/nemo_gym/train_data_utils.py b/nemo_gym/train_data_utils.py index 8a108b91df..8623031446 100644 --- a/nemo_gym/train_data_utils.py +++ b/nemo_gym/train_data_utils.py @@ -1,44 +1,40 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# 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 +from typing import Dict, List, Literal, Self, Optional, Union, Tuple + from abc import abstractmethod + from collections import Counter, defaultdict -from itertools import count, repeat -from pathlib import Path + +from itertools import repeat, count + +import json + from shutil import copyfileobj -from typing import Dict, List, Literal, Optional, Self, Union + +from pathlib import Path from devtools import pprint -from omegaconf import DictConfig -from pydantic import BaseModel, ConfigDict, Field, ValidationError + from tqdm.auto import tqdm -from nemo_gym.base_resources_server import BaseRunRequest +from pydantic import BaseModel, Field, ConfigDict, ValidationError + +from omegaconf import DictConfig + from nemo_gym.config_types import ( AGENT_REF_KEY, - AgentServerRef, + ServerInstanceConfig, DatasetConfig, DatasetType, DownloadJsonlDatasetGitlabConfig, - ServerInstanceConfig, + AgentServerRef, ) -from nemo_gym.gitlab_utils import download_jsonl_dataset from nemo_gym.global_config import ( GlobalConfigDictParser, GlobalConfigDictParserConfig, get_global_config_dict, ) +from nemo_gym.gitlab_utils import download_jsonl_dataset +from nemo_gym.base_resources_server import BaseRunRequest class TrainDataProcessorConfig(BaseModel): @@ -101,13 +97,19 @@ def _aggregate(self) -> Self: class DatasetMetrics(Accumulator): number_of_examples: int = Field(serialization_alias="Number of examples", default=0) - number_of_tools: AvgMinMax = Field(serialization_alias="Number of tools", default_factory=AvgMinMax) + number_of_tools: AvgMinMax = Field( + serialization_alias="Number of tools", default_factory=AvgMinMax + ) json_dumped_number_of_words: AvgMinMax = Field( serialization_alias="Json-dumped number of words (proxy for token count)", default_factory=AvgMinMax, ) - number_of_turns: AvgMinMax = Field(serialization_alias="Number of turns", default_factory=AvgMinMax) - temperature: AvgMinMax = Field(serialization_alias="Temperature", default_factory=AvgMinMax) + number_of_turns: AvgMinMax = Field( + serialization_alias="Number of turns", default_factory=AvgMinMax + ) + temperature: AvgMinMax = Field( + serialization_alias="Temperature", default_factory=AvgMinMax + ) # TODO: Number of unique create params, Number of unique user messages, other sampling params, etc @@ -128,6 +130,72 @@ def _aggregate(self: Self) -> Self: ) +def compute_sample_metrics(sample_dict_str: str) -> Tuple[DatasetMetrics, bool]: + try: + sample_dict = json.loads(sample_dict_str) + except json.JSONDecodeError: + return DatasetMetrics(), True + + try: + sample = BaseRunRequest.model_validate(sample_dict) + except ValidationError: + return DatasetMetrics(), True + + responses_create_params = sample.responses_create_params + responses_create_params = responses_create_params.model_dump(exclude_unset=True) + inputs = responses_create_params.get("input") + + number_of_tools_metrics = AvgMinMax() + if responses_create_params.get("tools") is not None: + number_of_tools = len(responses_create_params["tools"]) + number_of_tools_metrics = AvgMinMax( + total=1, + average=number_of_tools, + min=number_of_tools, + max=number_of_tools, + ) + + if isinstance(inputs, str): + inputs = [{"role": "user", "content": inputs}] + user_inputs = [i for i in inputs if i.get("role") == "user"] if inputs else [] + number_of_turns_metrics = AvgMinMax() + if user_inputs: + number_of_turns = len(user_inputs) + number_of_turns_metrics = AvgMinMax( + total=1, + average=number_of_turns, + min=number_of_turns, + max=number_of_turns, + ) + + temperature_metrics = AvgMinMax() + if responses_create_params.get("temperature") is not None: + temperature = responses_create_params["temperature"] + temperature_metrics = AvgMinMax( + total=1, + average=temperature, + min=temperature, + max=temperature, + ) + + json_dumped_number_of_words = len(json.dumps(responses_create_params).split()) + json_dumped_number_of_words_metrics = AvgMinMax( + total=1, + average=json_dumped_number_of_words, + min=json_dumped_number_of_words, + max=json_dumped_number_of_words, + ) + + metrics = DatasetMetrics( + number_of_examples=1, + number_of_tools=number_of_tools_metrics, + json_dumped_number_of_words=json_dumped_number_of_words_metrics, + number_of_turns=number_of_turns_metrics, + temperature=temperature_metrics, + ) + return metrics, False + + class DatasetValidatorState(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -144,7 +212,9 @@ def run(self, global_config_dict: DictConfig): # pragma: no cover config = TrainDataProcessorConfig.model_validate(global_config_dict) self._print_title("Load and validate server instance configs") - server_instance_configs = self.load_and_validate_server_instance_configs(config, global_config_dict) + server_instance_configs = self.load_and_validate_server_instance_configs( + config, global_config_dict + ) self._print_title( f"Load datasets. Missing datasets {'**WILL**' if config.should_download else 'will **NOT**'} be downloaded." @@ -152,10 +222,14 @@ def run(self, global_config_dict: DictConfig): # pragma: no cover self.load_datasets(config, server_instance_configs) self._print_title("Validate samples and aggregate metrics") - dataset_type_to_aggregate_metrics = self.validate_samples_and_aggregate_metrics(server_instance_configs) + dataset_type_to_aggregate_metrics = self.validate_samples_and_aggregate_metrics( + server_instance_configs + ) self._print_title("Collate samples and aggregate metrics") - self.collate_samples(config, server_instance_configs, dataset_type_to_aggregate_metrics) + self.collate_samples( + config, server_instance_configs, dataset_type_to_aggregate_metrics + ) self._print_title("Finished!") @@ -163,9 +237,9 @@ def _print_title(self, title: str) -> None: # pragma: no cover print(f""" {"#" * 100} -# +# # {title} -# +# {"#" * 100} """) @@ -173,13 +247,19 @@ def load_and_validate_server_instance_configs( self, config: TrainDataProcessorConfig, global_config_dict: DictConfig ) -> List[ServerInstanceConfig]: parser = GlobalConfigDictParser() - server_instance_configs = parser.filter_for_server_instance_configs(global_config_dict) + server_instance_configs = parser.filter_for_server_instance_configs( + global_config_dict + ) agent_configs: List[ServerInstanceConfig] = [ - c for c in server_instance_configs if c.SERVER_TYPE == "responses_api_agents" + c + for c in server_instance_configs + if c.SERVER_TYPE == "responses_api_agents" ] - server_names_list_str = "\n- ".join([""] + [f"{c.name} ({c.SERVER_TYPE})" for c in server_instance_configs]) + server_names_list_str = "\n- ".join( + [""] + [f"{c.name} ({c.SERVER_TYPE})" for c in server_instance_configs] + ) print( f"Found {len(server_instance_configs)} server instance configs ({len(agent_configs)} agent configs):{server_names_list_str}\n\n" ) @@ -192,7 +272,9 @@ def load_and_validate_server_instance_configs( else: agent_configs_without_data.append(agent_config) - server_names_list_str = "\n- ".join([""] + [f"{c.name} ({c.SERVER_TYPE})" for c in agent_configs_without_data]) + server_names_list_str = "\n- ".join( + [""] + [f"{c.name} ({c.SERVER_TYPE})" for c in agent_configs_without_data] + ) print( f"Found {len(agent_configs_without_data)} agent server instance configs withOUT datasets:{server_names_list_str}\n\n" ) @@ -200,7 +282,9 @@ def load_and_validate_server_instance_configs( server_names_list_str = "" for c in agent_configs_with_data: server_str = f"\n- {c.name}" - datasets_str = "\n - ".join([""] + [f"{d.name} ({d.type})" for d in c.datasets]) + datasets_str = "\n - ".join( + [""] + [f"{d.name} ({d.type})" for d in c.datasets] + ) server_names_list_str += f"{server_str}{datasets_str}" print( f"Found {len(agent_configs_with_data)} agent server instance configs WITH datasets:{server_names_list_str}\n\n" @@ -210,7 +294,9 @@ def load_and_validate_server_instance_configs( in_scope_dataset_types = config.in_scope_dataset_types agent_configs_with_in_scope_datasets: List[ServerInstanceConfig] = [] for agent_config in agent_configs_with_data: - in_scope_datasets = [d for d in agent_config.datasets if d.type in in_scope_dataset_types] + in_scope_datasets = [ + d for d in agent_config.datasets if d.type in in_scope_dataset_types + ] if not in_scope_datasets: continue @@ -221,9 +307,13 @@ def load_and_validate_server_instance_configs( server_names_list_str = "" for c in agent_configs_with_in_scope_datasets: server_str = f"\n- {c.name}" - datasets_str = "\n - ".join([""] + [f"{d.name} ({d.type})" for d in c.datasets]) + datasets_str = "\n - ".join( + [""] + [f"{d.name} ({d.type})" for d in c.datasets] + ) server_names_list_str += f"{server_str}{datasets_str}" - print(f"In scope dataset types for `{config.mode}` mode: {in_scope_dataset_types}") + print( + f"In scope dataset types for `{config.mode}` mode: {in_scope_dataset_types}" + ) print( f"Found {len(agent_configs_with_data)} agent server instance configs with in-scope datasets:{server_names_list_str}" ) @@ -248,18 +338,26 @@ def load_datasets( server_names_list_str = "" for server_name, datasets in local_datasets_found.items(): - datasets_str = "\n - ".join([""] + [f"{d.name} ({d.type})" for d in datasets]) + datasets_str = "\n - ".join( + [""] + [f"{d.name} ({d.type})" for d in datasets] + ) server_names_list_str += f"\n- {server_name}{datasets_str}" - print(f"FOUND the following datasets at their local paths:{server_names_list_str}\n\n") + print( + f"FOUND the following datasets at their local paths:{server_names_list_str}\n\n" + ) server_names_list_str = "" for server_name, datasets in local_datasets_not_found.items(): - datasets_str = "\n - ".join([""] + [f"{d.name} ({d.type})" for d in datasets]) + datasets_str = "\n - ".join( + [""] + [f"{d.name} ({d.type})" for d in datasets] + ) server_names_list_str += f"\n- {server_name}{datasets_str}" print(f"MISSING the following datasets:{server_names_list_str}\n\n") if config.mode == "example_validation": - assert not local_datasets_not_found, "You must provide the above missing example jsonl files!" + assert not local_datasets_not_found, ( + "You must provide the above missing example jsonl files!" + ) if not config.should_download: assert not local_datasets_not_found, ( "Missing local datasets. You must provide local datasets since download is disabled. Run with `+should_download=true` to enable downloading." @@ -273,7 +371,9 @@ def load_datasets( download_config = DownloadJsonlDatasetGitlabConfig.model_validate( d.gitlab_identifier.model_dump() | {"output_fpath": d.jsonl_fpath} ) - print(f"Downloading dataset `{d.name}` from `{server_name}` using {download_config}") + print( + f"Downloading dataset `{d.name}` from `{server_name}` using {download_config}" + ) download_jsonl_dataset(download_config) ######################################## @@ -283,72 +383,13 @@ def load_datasets( def _validate_samples_and_aggregate_metrics_single_sample( self, state: DatasetValidatorState, sample_idx: int, sample_dict_str: str ) -> None: - try: - sample_dict = json.loads(sample_dict_str) - except json.JSONDecodeError: - state.offending_example_idxs.append(sample_idx) - return - - try: - sample = BaseRunRequest.model_validate(sample_dict) - except ValidationError: + metrics, is_offending = compute_sample_metrics(sample_dict_str) + if is_offending: state.offending_example_idxs.append(sample_idx) return + sample_dict = json.loads(sample_dict_str) state.key_counts.update(sample_dict.keys()) - - responses_create_params = sample.responses_create_params - responses_create_params = responses_create_params.model_dump(exclude_unset=True) - inputs = responses_create_params["input"] - - number_of_tools_metrics = AvgMinMax() - if responses_create_params.get("tools") is not None: - number_of_tools = len(responses_create_params["tools"]) - number_of_tools_metrics = AvgMinMax( - total=1, - average=number_of_tools, - min=number_of_tools, - max=number_of_tools, - ) - - if isinstance(inputs, str): - inputs = [{"role": "user", "content": inputs}] - user_inputs = [i for i in inputs if i.get("role") == "user"] - number_of_turns_metrics = AvgMinMax() - if user_inputs: - number_of_turns = len(user_inputs) - number_of_turns_metrics = AvgMinMax( - total=1, - average=number_of_turns, - min=number_of_turns, - max=number_of_turns, - ) - - temperature_metrics = AvgMinMax() - if responses_create_params.get("temperature") is not None: - temperature = responses_create_params["temperature"] - temperature_metrics = AvgMinMax( - total=1, - average=temperature, - min=temperature, - max=temperature, - ) - - json_dumped_number_of_words = len(json.dumps(responses_create_params).split()) - json_dumped_number_of_words_metrics = AvgMinMax( - total=1, - average=json_dumped_number_of_words, - min=json_dumped_number_of_words, - max=json_dumped_number_of_words, - ) - - metrics = DatasetMetrics( - number_of_examples=1, - number_of_tools=number_of_tools_metrics, - json_dumped_number_of_words=json_dumped_number_of_words_metrics, - number_of_turns=number_of_turns_metrics, - temperature=temperature_metrics, - ) state.metrics.add(metrics) def _validate_samples_and_aggregate_metrics_single_dataset( @@ -368,7 +409,9 @@ def _validate_samples_and_aggregate_metrics_single_dataset( return state - def _validate_aggregate_metrics(self, aggregate_metrics_dict: Dict, metrics_fpath: Path) -> Optional[Path]: + def _validate_aggregate_metrics( + self, aggregate_metrics_dict: Dict, metrics_fpath: Path + ) -> Optional[Path]: """ Returns the conflicting metrics fpath if invalid. Else returns None """ @@ -376,7 +419,9 @@ def _validate_aggregate_metrics(self, aggregate_metrics_dict: Dict, metrics_fpat with open(metrics_fpath) as f: previous_aggregate_metrics_dict = json.load(f) if aggregate_metrics_dict != previous_aggregate_metrics_dict: - conflicting_metrics_fpath = metrics_fpath.with_name(f"{metrics_fpath.stem}_conflict.json") + conflicting_metrics_fpath = metrics_fpath.with_name( + f"{metrics_fpath.stem}_conflict.json" + ) with open(conflicting_metrics_fpath, "w") as f: json.dump(aggregate_metrics_dict, f, indent=4) @@ -386,7 +431,9 @@ def validate_samples_and_aggregate_metrics( self, server_instance_configs: List[ServerInstanceConfig] ) -> Dict[str, DatasetMetrics]: conflicting_fpaths: List[str] = [] - dataset_type_to_aggregate_metrics: Dict[str, DatasetMetrics] = defaultdict(DatasetMetrics) + dataset_type_to_aggregate_metrics: Dict[str, DatasetMetrics] = defaultdict( + DatasetMetrics + ) for c in server_instance_configs: for d in c.datasets: state = self._validate_samples_and_aggregate_metrics_single_dataset(d) @@ -395,7 +442,9 @@ def validate_samples_and_aggregate_metrics( aggregate_metrics = state.metrics.aggregate() - aggregate_metrics_dict = aggregate_metrics.model_dump(mode="json", by_alias=True) + aggregate_metrics_dict = aggregate_metrics.model_dump( + mode="json", by_alias=True + ) aggregate_metrics_dict = d.model_dump() | aggregate_metrics_dict data_fpath = Path(d.jsonl_fpath) @@ -416,7 +465,9 @@ def validate_samples_and_aggregate_metrics( if conflicting_fpaths: conflicting_fpaths_str = "\n- ".join([""] + conflicting_fpaths) - raise ValueError(f"Found conflicting aggregate metrics that need to be corrected:{conflicting_fpaths_str}") + raise ValueError( + f"Found conflicting aggregate metrics that need to be corrected:{conflicting_fpaths_str}" + ) return dict(dataset_type_to_aggregate_metrics) @@ -440,7 +491,9 @@ def _collate_samples_single_type( with open(data_path) as source, open(prepare_path, "w") as target: for line in tqdm(source, desc=f"Preparing data at {data_path}"): d = json.loads(line) - d[AGENT_REF_KEY] = AgentServerRef(type="responses_api_agents", name=c.name).model_dump() + d[AGENT_REF_KEY] = AgentServerRef( + type="responses_api_agents", name=c.name + ).model_dump() target.write(f"{json.dumps(d)}\n") paths_to_collate.append(prepare_path) @@ -462,7 +515,9 @@ def collate_samples( aggregate_metrics = dataset_type_to_aggregate_metrics[type] aggregate_metrics = aggregate_metrics.aggregate() - aggregate_metrics_dict = aggregate_metrics.model_dump(mode="json", by_alias=True) + aggregate_metrics_dict = aggregate_metrics.model_dump( + mode="json", by_alias=True + ) parent = Path(config.output_dirpath) parent.mkdir(exist_ok=True) @@ -495,9 +550,13 @@ def collate_samples( if conflicting_fpaths: conflicting_fpaths_str = "\n- ".join([""] + conflicting_fpaths) - raise ValueError(f"Found conflicting aggregate metrics that need to be corrected:{conflicting_fpaths_str}") + raise ValueError( + f"Found conflicting aggregate metrics that need to be corrected:{conflicting_fpaths_str}" + ) - final_fpaths_str = "\n- ".join([""] + [f"{type}: {fpath}" for type, fpath in final_fpaths.items()]) + final_fpaths_str = "\n- ".join( + [""] + [f"{type}: {fpath}" for type, fpath in final_fpaths.items()] + ) print(f"View your final data!{final_fpaths_str}") diff --git a/tests/nemo_gym/test_dataset_viewer.py b/tests/nemo_gym/test_dataset_viewer.py new file mode 100644 index 0000000000..eae5207aa6 --- /dev/null +++ b/tests/nemo_gym/test_dataset_viewer.py @@ -0,0 +1,115 @@ +from unittest.mock import patch, mock_open +from pytest import MonkeyPatch +from pydantic import BaseModel +from typing import Any + +import json + +from nemo_gym.dataset_viewer import ( + JsonlDatasetViewerConfig, + build_jsonl_dataset_viewer, + get_aggregate_metrics, +) + + +class TestDatasetViewer: + def test_sanity( + self, + ) -> None: + config = JsonlDatasetViewerConfig(jsonl_fpath="") + + # With tools + mock_content = r"""{"reward": 0.0, "accuracy": false, "set_overlap": 0.0, "original_term_minefield_hit": false, "order_instruction_following_failure": false, "id": 44, "expected_synonym_values": [489, 504], "expected_synonyms": ["Awake", "Alert"], "minefield_label": "Alive", "minefield_label_value": 497, "responses_create_params": {"input": [{"content": "# Instructions\nYou are an extraction agent. You will be provided a user query and you need to use the tools provided to you to extract list of synonym values. You will be provided with a bunch of synonyms for each. For each term, please see if it's relevant to the user query and get the values for each synonym as appropriate. You must get and extract the values for every synonym that appears in this list. Please output synonym values in the order they appear in the available synonyms below.\n\n# Available synonyms\nThe term 'Win' has a synonym 'Outperform'.\nThe term 'Empty' has a synonym 'Desolate'.\nThe term 'Open' has a synonym 'Revealed'.\nThe term 'Dry' has a synonym 'Arid'.\nThe term 'Bad' has a synonym 'Wicked'.\nThe term 'Retreat' has a synonym 'Draw back'.\nThe term 'Empty' has a synonym 'Vacant'.\nThe term 'Thick' has a synonym 'Fat'.\nThe term 'Stormy' has a synonym 'Gale-force'.\nThe term 'Rough' has a synonym 'Coarse'.\nThe term 'Day' has a synonym 'Afternoon'.\nThe term 'Quiet' has a synonym 'Hushed'.\nThe term 'Day' has a synonym 'Sunrise'.\nThe term 'Closed' has a synonym 'Covered'.\nThe term 'Early' has a synonym 'Preliminary'.\nThe term 'Night' has a synonym 'Midnight'.\nThe term 'Light' has a synonym 'Clear'.\nThe term 'Wide' has a synonym 'Comprehensive'.\nThe term 'Ugly' has a synonym 'Grotesque'.\nThe term 'Insult' has a synonym 'Belittle'.\nThe term 'Far' has a synonym 'Remote'.\nThe term 'Up' has a synonym 'Higher'.\nThe term 'Stormy' has a synonym 'Tempestuous'.\nThe term 'Dead' has a synonym 'Deceased'.\nThe term 'Dim' has a synonym 'Faint'.\nThe term 'Thick' has a synonym 'Heavy'.\nThe term 'Failure' has a synonym 'Loss'.\nThe term 'Sad' has a synonym 'Depressed'.\nThe term 'Thin' has a synonym 'Slender'.\nThe term 'Dry' has a synonym 'Dehydrated'.\nThe term 'Dirty' has a synonym 'Muddy'.\nThe term 'Fast' has a synonym 'Brisk'.\nThe term 'Defeat' has a synonym 'Failure'.\nThe term 'Sharp' has a synonym 'Tapered'.\nThe term 'Sharp' has a synonym 'Piercing'.\nThe term 'Rough' has a synonym 'Grainy'.\nThe term 'Cowardly' has a synonym 'Craven'.\nThe term 'False' has a synonym 'Incorrect'.\nThe term 'Sad' has a synonym 'Unhappy'.\nThe term 'Brave' has a synonym 'Stouthearted'.\nThe term 'Cowardly' has a synonym 'Yellow'.\nThe term 'Thin' has a synonym 'Skinny'.\nThe term 'Outside' has a synonym 'External'.\nThe term 'Wet' has a synonym 'Soaked'.\nThe term 'Sad' has a synonym 'Heartbroken'.\nThe term 'Success' has a synonym 'Victory'.\nThe term 'Cowardly' has a synonym 'Spineless'.\nThe term 'Tight' has a synonym 'Compact'.\nThe term 'Strong' has a synonym 'Muscular'.\nThe term 'Difficult' has a synonym 'Demanding'.\nThe term 'Old' has a synonym 'Historic'.\nThe term 'Rich' has a synonym 'Opulent'.\nThe term 'Far' has a synonym 'Away'.\nThe term 'Easy' has a synonym 'Effortless'.\nThe term 'Short' has a synonym 'Little'.\nThe term 'Win' has a synonym 'Achieve'.\nThe term 'Compliment' has a synonym 'Extol'.\nThe term 'Advance' has a synonym 'Rise'.\nThe term 'Soft' has a synonym 'Tender'.\nThe term 'Narrow' has a synonym 'Restrictive'.\nThe term 'Dark' has a synonym 'Dusky'.\nThe term 'High' has a synonym 'Prominent'.\nThe term 'Calm' has a synonym 'Undisturbed'.\nThe term 'Closed' has a synonym 'Locked'.\nThe term 'Compulsory' has a synonym 'Statutory'.\nThe term 'Alive' has a synonym 'Awake'.\nThe term 'Weak' has a synonym 'Powerless'.\nThe term 'Difficult' has a synonym 'Grueling'.\nThe term 'Reject' has a synonym 'Rebuff'.\nThe term 'Slow' has a synonym 'Leisurely'.\nThe term 'Clean' has a synonym 'Unsoiled'.\nThe term 'Compulsory' has a synonym 'Obligatory'.\nThe term 'Short' has a synonym 'Diminutive'.\nThe term 'Night' has a synonym 'Nightfall'.\nThe term 'Near' has a synonym 'Proximate'.\nThe term 'Ugly' has a synonym 'Homely'.\nThe term 'Wrong' has a synonym 'Amiss'.\nThe term 'Bad' has a synonym 'Terrible'.\nThe term 'Visible' has a synonym 'Noticeable'.\nThe term 'Near' has a synonym 'In proximity'.\nThe term 'Cold' has a synonym 'Frosty'.\nThe term 'Wrong' has a synonym 'False'.\nThe term 'Soft' has a synonym 'Velvety'.\nThe term 'Day' has a synonym 'Bright'.\nThe term 'Young' has a synonym 'Budding'.\nThe term 'Smelly' has a synonym 'Rancid'.\nThe term 'Low' has a synonym 'Diminished'.\nThe term 'Small' has a synonym 'Microscopic'.\nThe term 'Calm' has a synonym 'Unruffled'.\nThe term 'Empty' has a synonym 'Void'.\nThe term 'Open' has a synonym 'Available'.\nThe term 'Far' has a synonym 'Removed'.\nThe term 'Young' has a synonym 'New'.\nThe term 'Ascend' has a synonym 'Mount'.\nThe term 'Ugly' has a synonym 'Hideous'.\nThe term 'Weak' has a synonym 'Frail'.\nThe term 'Wet' has a synonym 'Damp'.\nThe term 'Tall' has a synonym 'Sky-high'.\nThe term 'Down' has a synonym 'Depressed'.\nThe term 'Happy' has a synonym 'Cheerful'.\nThe term 'Alive' has a synonym 'Alert'.\nThe term 'Easy' has a synonym 'Light'.\nThe term 'Accept' has a synonym 'Receive'.\nThe term 'Advance' has a synonym 'Headway'.\nThe term 'Dim' has a synonym 'Dull'.\nThe term 'Tall' has a synonym 'Towering'.\nThe term 'Fragrant' has a synonym 'Balmy'.\nThe term 'Happy' has a synonym 'Pleased'.\nThe term 'Down' has a synonym 'Drop'.\nThe term 'Hard' has a synonym 'Rigid'.\nThe term 'Loud' has a synonym 'Noisy'.\nThe term 'Light' has a synonym 'Shiny'.\nThe term 'Early' has a synonym 'Prior'.\nThe term 'Hot' has a synonym 'Blazing'.\nThe term 'Light (weight)' has a synonym 'Slim'.\nThe term 'Accept' has a synonym 'Acknowledge'.\nThe term 'Quiet' has a synonym 'Peaceful'.\nThe term 'Outside' has a synonym 'Outdoors'.\nThe term 'Easy' has a synonym 'Painless'.\nThe term 'Success' has a synonym 'Conquest'.\nThe term 'Hard' has a synonym 'Solid'.\nThe term 'Failure' has a synonym 'Setback'.\nThe term 'Low' has a synonym 'Short'.\nThe term 'Late' has a synonym 'Overdue'.\nThe term 'Wet' has a synonym 'Waterlogged'.\nThe term 'Strong' has a synonym 'Forceful'.\nThe term 'Hot' has a synonym 'Warm'.\nThe term 'Dark' has a synonym 'Tenebrous'.\nThe term 'Light (weight)' has a synonym 'Flimsy'.\nThe term 'Smelly' has a synonym 'Pungent'.\nThe term 'Soft' has a synonym 'Mild'.\nThe term 'Early' has a synonym 'First'.\nThe term 'Dirty' has a synonym 'Squalid'.\nThe term 'Dead' has a synonym 'Lifeless'.\nThe term 'Bitter' has a synonym 'Astringent'.\nThe term 'False' has a synonym 'Fallacious'.\nThe term 'Defeat' has a synonym 'Collapse'.\nThe term 'Loud' has a synonym 'Blaring'.\nThe term 'Dull' has a synonym 'Dim'.\nThe term 'Stormy' has a synonym 'Wild'.\nThe term 'Narrow' has a synonym 'Compressed'.\nThe term 'Rich' has a synonym 'Flush'.\nThe term 'Invisible' has a synonym 'Obscured'.\nThe term 'Slow' has a synonym 'Dragging'.\nThe term 'Young' has a synonym 'Juvenile'.\nThe term 'Bitter' has a synonym 'Caustic'.\nThe term 'Old' has a synonym 'Elderly'.\nThe term 'Slow' has a synonym 'Sluggish'.\nThe term 'Ascend' has a synonym 'Go up'.\nThe term 'Down' has a synonym 'Sink'.\nThe term 'Descend' has a synonym 'Subside'.\nThe term 'Small' has a synonym 'Little'.\nThe term 'High' has a synonym 'Soaring'.\nThe term 'Up' has a synonym 'Climb'.\nThe term 'Calm' has a synonym 'Relaxed'.\nThe term 'Rich' has a synonym 'Well-off'.\nThe term 'Light (weight)' has a synonym 'Breezy'.\nThe term 'Wrong' has a synonym 'Inaccurate'.\nThe term 'Dirty' has a synonym 'Soiled'.\nThe term 'Late' has a synonym 'Unpunctual'.\nThe term 'Quiet' has a synonym 'Low'.\nThe term 'Descend' has a synonym 'Dismount'.\nThe term 'Compliment' has a synonym 'Praise'.\nThe term 'Open' has a synonym 'Unsealed'.\nThe term 'Dull' has a synonym 'Blunt'.\nThe term 'Small' has a synonym 'Minor'.\nThe term 'Retreat' has a synonym 'Escape'.\nThe term 'Fast' has a synonym 'Hasty'.\nThe term 'Invisible' has a synonym 'Secret'.\nThe term 'Success' has a synonym 'Achievement'.\nThe term 'Retreat' has a synonym 'Fall back'.\nThe term 'Cold' has a synonym 'Icy'.\nThe term 'Hard' has a synonym 'Stiff'.\nThe term 'Insult' has a synonym 'Deride'.\nThe term 'Night' has a synonym 'Nocturne'.\nThe term 'Tight' has a synonym 'Firm'.\nThe term 'Accept' has a synonym 'Consent'.\nThe term 'Victory' has a synonym 'Supremacy'.\nThe term 'Old' has a synonym 'Vintage'.\nThe term 'Dry' has a synonym 'Desiccated'.\nThe term 'Narrow' has a synonym 'Pinched'.\nThe term 'Clean' has a synonym 'Sterile'.\nThe term 'Visible' has a synonym 'Perceptible'.\nThe term 'Victory' has a synonym 'Win'.\nThe term 'Advance' has a synonym 'Step up'.\nThe term 'Sharp' has a synonym 'Edged'.\nThe term 'Wide' has a synonym 'Outspread'.\nThe term 'Low' has a synonym 'Flat'.\nThe term 'Closed' has a synonym 'Fastened'.\nThe term 'False' has a synonym 'Untrue'.\nThe term 'Brave' has a synonym 'Bold'.\nThe term 'Reject' has a synonym 'Refuse'.\nThe term 'Fragrant' has a synonym 'Perfumed'.\n\n# Example\nFor example, if the user query is \"I'm very warm\", the term you should focus on is \"hot\". According to the synonyms above, the term \"hot\" has the synonyms \"Blazing\" and \"Warm\", in that order. You need to get synonym values for \"Blazing\" and \"Warm\", let's say those are 5 and 6 respectively, and extract the result of those synonym values i.e. [5, 6] with 5 (blazing) first then 6 (warm) since that is the order they appear in the list of synonyms above.", "role": "system"}, {"content": "How does the human body's response to danger highlight the instinct for survival?", "role": "user"}], "parallel_tool_calls": false, "tools": [{"name": "get_synonym_value", "parameters": {"properties": {"synonym": {"type": "string", "title": "Synonym", "description": "The synonym to get the value for."}}, "type": "object", "required": ["synonym"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Get the synonym value for a synonym.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"synonym_value\": {\"type\": \"integer\", \"title\": \"Synonym Value\", \"description\": \"The value for this synonym.\"}}, \"type\": \"object\", \"required\": [\"synonym_value\"]}\n"}, {"name": "extract_synonym_values", "parameters": {"properties": {"synonym_values": {"items": {"type": "integer"}, "type": "array", "title": "Synonym Values", "description": "The synonym values corresponding to the term for the user query."}}, "type": "object", "required": ["synonym_values"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Extract the synonym values you retrieved for the term that is relevant to the user query.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"success\": {\"type\": \"boolean\", \"title\": \"Success\", \"description\": \"Success.\"}}, \"type\": \"object\", \"required\": [\"success\"]}\n"}]}, "response": {"id": "resp_689038d64ad081929f6f36d2f2554431063ce0bbdad9d001", "created_at": 1754282198.0, "error": null, "incomplete_details": null, "instructions": null, "metadata": {}, "model": "gpt-4.1-2025-04-14", "object": "response", "output": [{"content": [{"annotations": [], "text": "fake chat message", "type": "output_text"}], "role": "assistant", "type": "message", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"summary": [{"type": "summary_text", "text": "fake reasoning"}], "type": "reasoning", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"arguments": "{\"synonym\":\"Survival\"}", "call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "name": "get_synonym_value", "type": "function_call", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "output": "{\"synonym_value\": 860}", "type": "function_call_output"}, {"arguments": "{\"synonym_values\":[860]}", "call_id": "call_N4Kr3NJJohoTaxSL5DkkG4W6", "name": "extract_synonym_values", "type": "function_call", "id": "fc_689038d6c71c8192b7ab0de7cc655d98063ce0bbdad9d001", "status": "completed"}], "parallel_tool_calls": false, "temperature": 1.0, "tool_choice": "auto", "tools": [{"name": "get_synonym_value", "parameters": {"properties": {"synonym": {"type": "string", "title": "Synonym", "description": "The synonym to get the value for."}}, "type": "object", "required": ["synonym"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Get the synonym value for a synonym.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"synonym_value\": {\"type\": \"integer\", \"title\": \"Synonym Value\", \"description\": \"The value for this synonym.\"}}, \"type\": \"object\", \"required\": [\"synonym_value\"]}\n"}, {"name": "extract_synonym_values", "parameters": {"properties": {"synonym_values": {"items": {"type": "integer"}, "type": "array", "title": "Synonym Values", "description": "The synonym values corresponding to the term for the user query."}}, "type": "object", "required": ["synonym_values"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Extract the synonym values you retrieved for the term that is relevant to the user query.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"success\": {\"type\": \"boolean\", \"title\": \"Success\", \"description\": \"Success.\"}}, \"type\": \"object\", \"required\": [\"success\"]}\n"}], "top_p": 1.0, "background": false, "max_output_tokens": null, "max_tool_calls": null, "previous_response_id": null, "prompt": null, "reasoning": {"effort": null, "generate_summary": null, "summary": null}, "service_tier": "default", "status": "completed", "text": {"format": {"type": "text"}}, "top_logprobs": 0, "truncation": "disabled", "usage": {"input_tokens": 2864, "input_tokens_details": {"cached_tokens": 2798}, "output_tokens": 19, "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": 2883}, "user": null, "prompt_cache_key": null, "safety_identifier": null, "store": true}}""" + with patch("builtins.open", mock_open(read_data=mock_content)): + build_jsonl_dataset_viewer(config) + + # Without tools in create params. Responses will always have tools + mock_content = r"""{"reward": 0.0, "accuracy": false, "set_overlap": 0.0, "original_term_minefield_hit": false, "order_instruction_following_failure": false, "id": 44, "expected_synonym_values": [489, 504], "expected_synonyms": ["Awake", "Alert"], "minefield_label": "Alive", "minefield_label_value": 497, "responses_create_params": {"input": [{"content": "# Instructions", "role": "system"}, {"content": "How does the human body's response to danger highlight the instinct for survival?", "role": "user"}]}, "response": {"id": "resp_689038d64ad081929f6f36d2f2554431063ce0bbdad9d001", "created_at": 1754282198.0, "error": null, "incomplete_details": null, "instructions": null, "metadata": {}, "model": "gpt-4.1-2025-04-14", "object": "response", "output": [{"content": [{"annotations": [], "text": "fake chat message", "type": "output_text"}], "role": "assistant", "type": "message", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"summary": [{"type": "summary_text", "text": "fake reasoning"}], "type": "reasoning", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"arguments": "{\"synonym\":\"Survival\"}", "call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "name": "get_synonym_value", "type": "function_call", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "output": "{\"synonym_value\": 860}", "type": "function_call_output"}, {"arguments": "{\"synonym_values\":[860]}", "call_id": "call_N4Kr3NJJohoTaxSL5DkkG4W6", "name": "extract_synonym_values", "type": "function_call", "id": "fc_689038d6c71c8192b7ab0de7cc655d98063ce0bbdad9d001", "status": "completed"}], "parallel_tool_calls": false, "temperature": 1.0, "tool_choice": "auto", "tools": [{"name": "get_synonym_value", "parameters": {"properties": {"synonym": {"type": "string", "title": "Synonym", "description": "The synonym to get the value for."}}, "type": "object", "required": ["synonym"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Get the synonym value for a synonym.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"synonym_value\": {\"type\": \"integer\", \"title\": \"Synonym Value\", \"description\": \"The value for this synonym.\"}}, \"type\": \"object\", \"required\": [\"synonym_value\"]}\n"}, {"name": "extract_synonym_values", "parameters": {"properties": {"synonym_values": {"items": {"type": "integer"}, "type": "array", "title": "Synonym Values", "description": "The synonym values corresponding to the term for the user query."}}, "type": "object", "required": ["synonym_values"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Extract the synonym values you retrieved for the term that is relevant to the user query.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"success\": {\"type\": \"boolean\", \"title\": \"Success\", \"description\": \"Success.\"}}, \"type\": \"object\", \"required\": [\"success\"]}\n"}], "top_p": 1.0, "background": false, "max_output_tokens": null, "max_tool_calls": null, "previous_response_id": null, "prompt": null, "reasoning": {"effort": null, "generate_summary": null, "summary": null}, "service_tier": "default", "status": "completed", "text": {"format": {"type": "text"}}, "top_logprobs": 0, "truncation": "disabled", "usage": {"input_tokens": 2864, "input_tokens_details": {"cached_tokens": 2798}, "output_tokens": 19, "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": 2883}, "user": null, "prompt_cache_key": null, "safety_identifier": null, "store": true}}""" + with patch("builtins.open", mock_open(read_data=mock_content)): + build_jsonl_dataset_viewer(config) + + def test_get_aggregate_metrics(self, monkeypatch: MonkeyPatch): + class DummySample(BaseModel): + responses_create_params: dict = {} + response: dict = {} + reward: float = 1.0 + accuracy: bool = True + set_overlap: float = 0.5 + unrelated_list: list = [] + unrelated_dict: dict = {} + + class DummySampleWithStrings(DummySample): + some_string: str + + config = JsonlDatasetViewerConfig(jsonl_fpath="") + samples = [ + DummySample(reward=1.0, accuracy=True, set_overlap=0.5), + DummySample(reward=0.0, accuracy=False, set_overlap=0.0), + DummySample(reward=0.5, accuracy=True, set_overlap=1.0), + ] + + samples_with_strings = [ + DummySampleWithStrings(reward=1.0, accuracy=True, some_string="asdf"), + DummySampleWithStrings(reward=0.0, accuracy=False, some_string="asdf"), + DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word1"), + DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word1"), + DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word2"), + ] + + def mock_compute_sample_metrics(line: str): + metrics = json.loads(line) + return metrics, False + + class DummyAgg: + def model_dump(self, by_alias=True): + return {} + + class DummyDatasetMetrics: + def add(self, metrics: Any): + pass + + def aggregate(self): + return DummyAgg() + + monkeypatch.setattr( + "nemo_gym.train_data_utils.compute_sample_metrics", + mock_compute_sample_metrics, + ) + monkeypatch.setattr( + "nemo_gym.train_data_utils.DatasetMetrics", DummyDatasetMetrics + ) + + with patch("builtins.open", mock_open(read_data="{}\n")): + result_1 = get_aggregate_metrics(config, samples) + + assert "reward" in result_1 + assert "accuracy" in result_1 + assert "set_overlap" in result_1 + + assert "unrelated_str" not in result_1 + assert "unrelated_list" not in result_1 + assert "unrelated_dict" not in result_1 + + assert "responses_create_params" not in result_1 + assert "response" not in result_1 + + # Check computed values + reward_stats = result_1["reward"] + assert reward_stats["Total # non-null values"] == 3 + assert reward_stats["Average"] == (1.0 + 0.0 + 0.5) / 3 + assert reward_stats["Min"] == 0.0 + assert reward_stats["Max"] == 1.0 + + # Check computed values with bools converted to int + accuracy_stats = result_1["accuracy"] + assert accuracy_stats["Total # non-null values"] == 3 + assert accuracy_stats["Average"] == (1 + 0 + 1) / 3 + assert accuracy_stats["Min"] == 0 + assert accuracy_stats["Max"] == 1 + + with patch("builtins.open", mock_open(read_data="{}\n")): + result_2 = get_aggregate_metrics(config, samples_with_strings) + + assert "some_string" in result_2 + assert result_2["some_string"]["unique_count"] == 3 + assert result_2["some_string"]["total_count"] == 5 From 65a7de126ffbf850ec613ccdd79d43f18ad8b910 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Fri, 5 Sep 2025 14:50:46 -0700 Subject: [PATCH 2/5] Restore copyright/license, move test Signed-off-by: Frankie Siino --- nemo_gym/dataset_viewer.py | 13 +++ nemo_gym/train_data_utils.py | 13 +++ tests/nemo_gym/test_dataset_viewer.py | 115 ------------------------ tests/unit_tests/test_dataset_viewer.py | 96 +++++++++++++++++++- 4 files changed, 120 insertions(+), 117 deletions(-) delete mode 100644 tests/nemo_gym/test_dataset_viewer.py diff --git a/nemo_gym/dataset_viewer.py b/nemo_gym/dataset_viewer.py index 93f229a127..8faa738c3f 100644 --- a/nemo_gym/dataset_viewer.py +++ b/nemo_gym/dataset_viewer.py @@ -1,3 +1,16 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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 typing import List, Dict, Any import json diff --git a/nemo_gym/train_data_utils.py b/nemo_gym/train_data_utils.py index 8623031446..16536520e5 100644 --- a/nemo_gym/train_data_utils.py +++ b/nemo_gym/train_data_utils.py @@ -1,3 +1,16 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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 typing import Dict, List, Literal, Self, Optional, Union, Tuple from abc import abstractmethod diff --git a/tests/nemo_gym/test_dataset_viewer.py b/tests/nemo_gym/test_dataset_viewer.py deleted file mode 100644 index eae5207aa6..0000000000 --- a/tests/nemo_gym/test_dataset_viewer.py +++ /dev/null @@ -1,115 +0,0 @@ -from unittest.mock import patch, mock_open -from pytest import MonkeyPatch -from pydantic import BaseModel -from typing import Any - -import json - -from nemo_gym.dataset_viewer import ( - JsonlDatasetViewerConfig, - build_jsonl_dataset_viewer, - get_aggregate_metrics, -) - - -class TestDatasetViewer: - def test_sanity( - self, - ) -> None: - config = JsonlDatasetViewerConfig(jsonl_fpath="") - - # With tools - mock_content = r"""{"reward": 0.0, "accuracy": false, "set_overlap": 0.0, "original_term_minefield_hit": false, "order_instruction_following_failure": false, "id": 44, "expected_synonym_values": [489, 504], "expected_synonyms": ["Awake", "Alert"], "minefield_label": "Alive", "minefield_label_value": 497, "responses_create_params": {"input": [{"content": "# Instructions\nYou are an extraction agent. You will be provided a user query and you need to use the tools provided to you to extract list of synonym values. You will be provided with a bunch of synonyms for each. For each term, please see if it's relevant to the user query and get the values for each synonym as appropriate. You must get and extract the values for every synonym that appears in this list. Please output synonym values in the order they appear in the available synonyms below.\n\n# Available synonyms\nThe term 'Win' has a synonym 'Outperform'.\nThe term 'Empty' has a synonym 'Desolate'.\nThe term 'Open' has a synonym 'Revealed'.\nThe term 'Dry' has a synonym 'Arid'.\nThe term 'Bad' has a synonym 'Wicked'.\nThe term 'Retreat' has a synonym 'Draw back'.\nThe term 'Empty' has a synonym 'Vacant'.\nThe term 'Thick' has a synonym 'Fat'.\nThe term 'Stormy' has a synonym 'Gale-force'.\nThe term 'Rough' has a synonym 'Coarse'.\nThe term 'Day' has a synonym 'Afternoon'.\nThe term 'Quiet' has a synonym 'Hushed'.\nThe term 'Day' has a synonym 'Sunrise'.\nThe term 'Closed' has a synonym 'Covered'.\nThe term 'Early' has a synonym 'Preliminary'.\nThe term 'Night' has a synonym 'Midnight'.\nThe term 'Light' has a synonym 'Clear'.\nThe term 'Wide' has a synonym 'Comprehensive'.\nThe term 'Ugly' has a synonym 'Grotesque'.\nThe term 'Insult' has a synonym 'Belittle'.\nThe term 'Far' has a synonym 'Remote'.\nThe term 'Up' has a synonym 'Higher'.\nThe term 'Stormy' has a synonym 'Tempestuous'.\nThe term 'Dead' has a synonym 'Deceased'.\nThe term 'Dim' has a synonym 'Faint'.\nThe term 'Thick' has a synonym 'Heavy'.\nThe term 'Failure' has a synonym 'Loss'.\nThe term 'Sad' has a synonym 'Depressed'.\nThe term 'Thin' has a synonym 'Slender'.\nThe term 'Dry' has a synonym 'Dehydrated'.\nThe term 'Dirty' has a synonym 'Muddy'.\nThe term 'Fast' has a synonym 'Brisk'.\nThe term 'Defeat' has a synonym 'Failure'.\nThe term 'Sharp' has a synonym 'Tapered'.\nThe term 'Sharp' has a synonym 'Piercing'.\nThe term 'Rough' has a synonym 'Grainy'.\nThe term 'Cowardly' has a synonym 'Craven'.\nThe term 'False' has a synonym 'Incorrect'.\nThe term 'Sad' has a synonym 'Unhappy'.\nThe term 'Brave' has a synonym 'Stouthearted'.\nThe term 'Cowardly' has a synonym 'Yellow'.\nThe term 'Thin' has a synonym 'Skinny'.\nThe term 'Outside' has a synonym 'External'.\nThe term 'Wet' has a synonym 'Soaked'.\nThe term 'Sad' has a synonym 'Heartbroken'.\nThe term 'Success' has a synonym 'Victory'.\nThe term 'Cowardly' has a synonym 'Spineless'.\nThe term 'Tight' has a synonym 'Compact'.\nThe term 'Strong' has a synonym 'Muscular'.\nThe term 'Difficult' has a synonym 'Demanding'.\nThe term 'Old' has a synonym 'Historic'.\nThe term 'Rich' has a synonym 'Opulent'.\nThe term 'Far' has a synonym 'Away'.\nThe term 'Easy' has a synonym 'Effortless'.\nThe term 'Short' has a synonym 'Little'.\nThe term 'Win' has a synonym 'Achieve'.\nThe term 'Compliment' has a synonym 'Extol'.\nThe term 'Advance' has a synonym 'Rise'.\nThe term 'Soft' has a synonym 'Tender'.\nThe term 'Narrow' has a synonym 'Restrictive'.\nThe term 'Dark' has a synonym 'Dusky'.\nThe term 'High' has a synonym 'Prominent'.\nThe term 'Calm' has a synonym 'Undisturbed'.\nThe term 'Closed' has a synonym 'Locked'.\nThe term 'Compulsory' has a synonym 'Statutory'.\nThe term 'Alive' has a synonym 'Awake'.\nThe term 'Weak' has a synonym 'Powerless'.\nThe term 'Difficult' has a synonym 'Grueling'.\nThe term 'Reject' has a synonym 'Rebuff'.\nThe term 'Slow' has a synonym 'Leisurely'.\nThe term 'Clean' has a synonym 'Unsoiled'.\nThe term 'Compulsory' has a synonym 'Obligatory'.\nThe term 'Short' has a synonym 'Diminutive'.\nThe term 'Night' has a synonym 'Nightfall'.\nThe term 'Near' has a synonym 'Proximate'.\nThe term 'Ugly' has a synonym 'Homely'.\nThe term 'Wrong' has a synonym 'Amiss'.\nThe term 'Bad' has a synonym 'Terrible'.\nThe term 'Visible' has a synonym 'Noticeable'.\nThe term 'Near' has a synonym 'In proximity'.\nThe term 'Cold' has a synonym 'Frosty'.\nThe term 'Wrong' has a synonym 'False'.\nThe term 'Soft' has a synonym 'Velvety'.\nThe term 'Day' has a synonym 'Bright'.\nThe term 'Young' has a synonym 'Budding'.\nThe term 'Smelly' has a synonym 'Rancid'.\nThe term 'Low' has a synonym 'Diminished'.\nThe term 'Small' has a synonym 'Microscopic'.\nThe term 'Calm' has a synonym 'Unruffled'.\nThe term 'Empty' has a synonym 'Void'.\nThe term 'Open' has a synonym 'Available'.\nThe term 'Far' has a synonym 'Removed'.\nThe term 'Young' has a synonym 'New'.\nThe term 'Ascend' has a synonym 'Mount'.\nThe term 'Ugly' has a synonym 'Hideous'.\nThe term 'Weak' has a synonym 'Frail'.\nThe term 'Wet' has a synonym 'Damp'.\nThe term 'Tall' has a synonym 'Sky-high'.\nThe term 'Down' has a synonym 'Depressed'.\nThe term 'Happy' has a synonym 'Cheerful'.\nThe term 'Alive' has a synonym 'Alert'.\nThe term 'Easy' has a synonym 'Light'.\nThe term 'Accept' has a synonym 'Receive'.\nThe term 'Advance' has a synonym 'Headway'.\nThe term 'Dim' has a synonym 'Dull'.\nThe term 'Tall' has a synonym 'Towering'.\nThe term 'Fragrant' has a synonym 'Balmy'.\nThe term 'Happy' has a synonym 'Pleased'.\nThe term 'Down' has a synonym 'Drop'.\nThe term 'Hard' has a synonym 'Rigid'.\nThe term 'Loud' has a synonym 'Noisy'.\nThe term 'Light' has a synonym 'Shiny'.\nThe term 'Early' has a synonym 'Prior'.\nThe term 'Hot' has a synonym 'Blazing'.\nThe term 'Light (weight)' has a synonym 'Slim'.\nThe term 'Accept' has a synonym 'Acknowledge'.\nThe term 'Quiet' has a synonym 'Peaceful'.\nThe term 'Outside' has a synonym 'Outdoors'.\nThe term 'Easy' has a synonym 'Painless'.\nThe term 'Success' has a synonym 'Conquest'.\nThe term 'Hard' has a synonym 'Solid'.\nThe term 'Failure' has a synonym 'Setback'.\nThe term 'Low' has a synonym 'Short'.\nThe term 'Late' has a synonym 'Overdue'.\nThe term 'Wet' has a synonym 'Waterlogged'.\nThe term 'Strong' has a synonym 'Forceful'.\nThe term 'Hot' has a synonym 'Warm'.\nThe term 'Dark' has a synonym 'Tenebrous'.\nThe term 'Light (weight)' has a synonym 'Flimsy'.\nThe term 'Smelly' has a synonym 'Pungent'.\nThe term 'Soft' has a synonym 'Mild'.\nThe term 'Early' has a synonym 'First'.\nThe term 'Dirty' has a synonym 'Squalid'.\nThe term 'Dead' has a synonym 'Lifeless'.\nThe term 'Bitter' has a synonym 'Astringent'.\nThe term 'False' has a synonym 'Fallacious'.\nThe term 'Defeat' has a synonym 'Collapse'.\nThe term 'Loud' has a synonym 'Blaring'.\nThe term 'Dull' has a synonym 'Dim'.\nThe term 'Stormy' has a synonym 'Wild'.\nThe term 'Narrow' has a synonym 'Compressed'.\nThe term 'Rich' has a synonym 'Flush'.\nThe term 'Invisible' has a synonym 'Obscured'.\nThe term 'Slow' has a synonym 'Dragging'.\nThe term 'Young' has a synonym 'Juvenile'.\nThe term 'Bitter' has a synonym 'Caustic'.\nThe term 'Old' has a synonym 'Elderly'.\nThe term 'Slow' has a synonym 'Sluggish'.\nThe term 'Ascend' has a synonym 'Go up'.\nThe term 'Down' has a synonym 'Sink'.\nThe term 'Descend' has a synonym 'Subside'.\nThe term 'Small' has a synonym 'Little'.\nThe term 'High' has a synonym 'Soaring'.\nThe term 'Up' has a synonym 'Climb'.\nThe term 'Calm' has a synonym 'Relaxed'.\nThe term 'Rich' has a synonym 'Well-off'.\nThe term 'Light (weight)' has a synonym 'Breezy'.\nThe term 'Wrong' has a synonym 'Inaccurate'.\nThe term 'Dirty' has a synonym 'Soiled'.\nThe term 'Late' has a synonym 'Unpunctual'.\nThe term 'Quiet' has a synonym 'Low'.\nThe term 'Descend' has a synonym 'Dismount'.\nThe term 'Compliment' has a synonym 'Praise'.\nThe term 'Open' has a synonym 'Unsealed'.\nThe term 'Dull' has a synonym 'Blunt'.\nThe term 'Small' has a synonym 'Minor'.\nThe term 'Retreat' has a synonym 'Escape'.\nThe term 'Fast' has a synonym 'Hasty'.\nThe term 'Invisible' has a synonym 'Secret'.\nThe term 'Success' has a synonym 'Achievement'.\nThe term 'Retreat' has a synonym 'Fall back'.\nThe term 'Cold' has a synonym 'Icy'.\nThe term 'Hard' has a synonym 'Stiff'.\nThe term 'Insult' has a synonym 'Deride'.\nThe term 'Night' has a synonym 'Nocturne'.\nThe term 'Tight' has a synonym 'Firm'.\nThe term 'Accept' has a synonym 'Consent'.\nThe term 'Victory' has a synonym 'Supremacy'.\nThe term 'Old' has a synonym 'Vintage'.\nThe term 'Dry' has a synonym 'Desiccated'.\nThe term 'Narrow' has a synonym 'Pinched'.\nThe term 'Clean' has a synonym 'Sterile'.\nThe term 'Visible' has a synonym 'Perceptible'.\nThe term 'Victory' has a synonym 'Win'.\nThe term 'Advance' has a synonym 'Step up'.\nThe term 'Sharp' has a synonym 'Edged'.\nThe term 'Wide' has a synonym 'Outspread'.\nThe term 'Low' has a synonym 'Flat'.\nThe term 'Closed' has a synonym 'Fastened'.\nThe term 'False' has a synonym 'Untrue'.\nThe term 'Brave' has a synonym 'Bold'.\nThe term 'Reject' has a synonym 'Refuse'.\nThe term 'Fragrant' has a synonym 'Perfumed'.\n\n# Example\nFor example, if the user query is \"I'm very warm\", the term you should focus on is \"hot\". According to the synonyms above, the term \"hot\" has the synonyms \"Blazing\" and \"Warm\", in that order. You need to get synonym values for \"Blazing\" and \"Warm\", let's say those are 5 and 6 respectively, and extract the result of those synonym values i.e. [5, 6] with 5 (blazing) first then 6 (warm) since that is the order they appear in the list of synonyms above.", "role": "system"}, {"content": "How does the human body's response to danger highlight the instinct for survival?", "role": "user"}], "parallel_tool_calls": false, "tools": [{"name": "get_synonym_value", "parameters": {"properties": {"synonym": {"type": "string", "title": "Synonym", "description": "The synonym to get the value for."}}, "type": "object", "required": ["synonym"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Get the synonym value for a synonym.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"synonym_value\": {\"type\": \"integer\", \"title\": \"Synonym Value\", \"description\": \"The value for this synonym.\"}}, \"type\": \"object\", \"required\": [\"synonym_value\"]}\n"}, {"name": "extract_synonym_values", "parameters": {"properties": {"synonym_values": {"items": {"type": "integer"}, "type": "array", "title": "Synonym Values", "description": "The synonym values corresponding to the term for the user query."}}, "type": "object", "required": ["synonym_values"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Extract the synonym values you retrieved for the term that is relevant to the user query.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"success\": {\"type\": \"boolean\", \"title\": \"Success\", \"description\": \"Success.\"}}, \"type\": \"object\", \"required\": [\"success\"]}\n"}]}, "response": {"id": "resp_689038d64ad081929f6f36d2f2554431063ce0bbdad9d001", "created_at": 1754282198.0, "error": null, "incomplete_details": null, "instructions": null, "metadata": {}, "model": "gpt-4.1-2025-04-14", "object": "response", "output": [{"content": [{"annotations": [], "text": "fake chat message", "type": "output_text"}], "role": "assistant", "type": "message", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"summary": [{"type": "summary_text", "text": "fake reasoning"}], "type": "reasoning", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"arguments": "{\"synonym\":\"Survival\"}", "call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "name": "get_synonym_value", "type": "function_call", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "output": "{\"synonym_value\": 860}", "type": "function_call_output"}, {"arguments": "{\"synonym_values\":[860]}", "call_id": "call_N4Kr3NJJohoTaxSL5DkkG4W6", "name": "extract_synonym_values", "type": "function_call", "id": "fc_689038d6c71c8192b7ab0de7cc655d98063ce0bbdad9d001", "status": "completed"}], "parallel_tool_calls": false, "temperature": 1.0, "tool_choice": "auto", "tools": [{"name": "get_synonym_value", "parameters": {"properties": {"synonym": {"type": "string", "title": "Synonym", "description": "The synonym to get the value for."}}, "type": "object", "required": ["synonym"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Get the synonym value for a synonym.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"synonym_value\": {\"type\": \"integer\", \"title\": \"Synonym Value\", \"description\": \"The value for this synonym.\"}}, \"type\": \"object\", \"required\": [\"synonym_value\"]}\n"}, {"name": "extract_synonym_values", "parameters": {"properties": {"synonym_values": {"items": {"type": "integer"}, "type": "array", "title": "Synonym Values", "description": "The synonym values corresponding to the term for the user query."}}, "type": "object", "required": ["synonym_values"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Extract the synonym values you retrieved for the term that is relevant to the user query.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"success\": {\"type\": \"boolean\", \"title\": \"Success\", \"description\": \"Success.\"}}, \"type\": \"object\", \"required\": [\"success\"]}\n"}], "top_p": 1.0, "background": false, "max_output_tokens": null, "max_tool_calls": null, "previous_response_id": null, "prompt": null, "reasoning": {"effort": null, "generate_summary": null, "summary": null}, "service_tier": "default", "status": "completed", "text": {"format": {"type": "text"}}, "top_logprobs": 0, "truncation": "disabled", "usage": {"input_tokens": 2864, "input_tokens_details": {"cached_tokens": 2798}, "output_tokens": 19, "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": 2883}, "user": null, "prompt_cache_key": null, "safety_identifier": null, "store": true}}""" - with patch("builtins.open", mock_open(read_data=mock_content)): - build_jsonl_dataset_viewer(config) - - # Without tools in create params. Responses will always have tools - mock_content = r"""{"reward": 0.0, "accuracy": false, "set_overlap": 0.0, "original_term_minefield_hit": false, "order_instruction_following_failure": false, "id": 44, "expected_synonym_values": [489, 504], "expected_synonyms": ["Awake", "Alert"], "minefield_label": "Alive", "minefield_label_value": 497, "responses_create_params": {"input": [{"content": "# Instructions", "role": "system"}, {"content": "How does the human body's response to danger highlight the instinct for survival?", "role": "user"}]}, "response": {"id": "resp_689038d64ad081929f6f36d2f2554431063ce0bbdad9d001", "created_at": 1754282198.0, "error": null, "incomplete_details": null, "instructions": null, "metadata": {}, "model": "gpt-4.1-2025-04-14", "object": "response", "output": [{"content": [{"annotations": [], "text": "fake chat message", "type": "output_text"}], "role": "assistant", "type": "message", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"summary": [{"type": "summary_text", "text": "fake reasoning"}], "type": "reasoning", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"arguments": "{\"synonym\":\"Survival\"}", "call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "name": "get_synonym_value", "type": "function_call", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "output": "{\"synonym_value\": 860}", "type": "function_call_output"}, {"arguments": "{\"synonym_values\":[860]}", "call_id": "call_N4Kr3NJJohoTaxSL5DkkG4W6", "name": "extract_synonym_values", "type": "function_call", "id": "fc_689038d6c71c8192b7ab0de7cc655d98063ce0bbdad9d001", "status": "completed"}], "parallel_tool_calls": false, "temperature": 1.0, "tool_choice": "auto", "tools": [{"name": "get_synonym_value", "parameters": {"properties": {"synonym": {"type": "string", "title": "Synonym", "description": "The synonym to get the value for."}}, "type": "object", "required": ["synonym"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Get the synonym value for a synonym.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"synonym_value\": {\"type\": \"integer\", \"title\": \"Synonym Value\", \"description\": \"The value for this synonym.\"}}, \"type\": \"object\", \"required\": [\"synonym_value\"]}\n"}, {"name": "extract_synonym_values", "parameters": {"properties": {"synonym_values": {"items": {"type": "integer"}, "type": "array", "title": "Synonym Values", "description": "The synonym values corresponding to the term for the user query."}}, "type": "object", "required": ["synonym_values"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Extract the synonym values you retrieved for the term that is relevant to the user query.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"success\": {\"type\": \"boolean\", \"title\": \"Success\", \"description\": \"Success.\"}}, \"type\": \"object\", \"required\": [\"success\"]}\n"}], "top_p": 1.0, "background": false, "max_output_tokens": null, "max_tool_calls": null, "previous_response_id": null, "prompt": null, "reasoning": {"effort": null, "generate_summary": null, "summary": null}, "service_tier": "default", "status": "completed", "text": {"format": {"type": "text"}}, "top_logprobs": 0, "truncation": "disabled", "usage": {"input_tokens": 2864, "input_tokens_details": {"cached_tokens": 2798}, "output_tokens": 19, "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": 2883}, "user": null, "prompt_cache_key": null, "safety_identifier": null, "store": true}}""" - with patch("builtins.open", mock_open(read_data=mock_content)): - build_jsonl_dataset_viewer(config) - - def test_get_aggregate_metrics(self, monkeypatch: MonkeyPatch): - class DummySample(BaseModel): - responses_create_params: dict = {} - response: dict = {} - reward: float = 1.0 - accuracy: bool = True - set_overlap: float = 0.5 - unrelated_list: list = [] - unrelated_dict: dict = {} - - class DummySampleWithStrings(DummySample): - some_string: str - - config = JsonlDatasetViewerConfig(jsonl_fpath="") - samples = [ - DummySample(reward=1.0, accuracy=True, set_overlap=0.5), - DummySample(reward=0.0, accuracy=False, set_overlap=0.0), - DummySample(reward=0.5, accuracy=True, set_overlap=1.0), - ] - - samples_with_strings = [ - DummySampleWithStrings(reward=1.0, accuracy=True, some_string="asdf"), - DummySampleWithStrings(reward=0.0, accuracy=False, some_string="asdf"), - DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word1"), - DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word1"), - DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word2"), - ] - - def mock_compute_sample_metrics(line: str): - metrics = json.loads(line) - return metrics, False - - class DummyAgg: - def model_dump(self, by_alias=True): - return {} - - class DummyDatasetMetrics: - def add(self, metrics: Any): - pass - - def aggregate(self): - return DummyAgg() - - monkeypatch.setattr( - "nemo_gym.train_data_utils.compute_sample_metrics", - mock_compute_sample_metrics, - ) - monkeypatch.setattr( - "nemo_gym.train_data_utils.DatasetMetrics", DummyDatasetMetrics - ) - - with patch("builtins.open", mock_open(read_data="{}\n")): - result_1 = get_aggregate_metrics(config, samples) - - assert "reward" in result_1 - assert "accuracy" in result_1 - assert "set_overlap" in result_1 - - assert "unrelated_str" not in result_1 - assert "unrelated_list" not in result_1 - assert "unrelated_dict" not in result_1 - - assert "responses_create_params" not in result_1 - assert "response" not in result_1 - - # Check computed values - reward_stats = result_1["reward"] - assert reward_stats["Total # non-null values"] == 3 - assert reward_stats["Average"] == (1.0 + 0.0 + 0.5) / 3 - assert reward_stats["Min"] == 0.0 - assert reward_stats["Max"] == 1.0 - - # Check computed values with bools converted to int - accuracy_stats = result_1["accuracy"] - assert accuracy_stats["Total # non-null values"] == 3 - assert accuracy_stats["Average"] == (1 + 0 + 1) / 3 - assert accuracy_stats["Min"] == 0 - assert accuracy_stats["Max"] == 1 - - with patch("builtins.open", mock_open(read_data="{}\n")): - result_2 = get_aggregate_metrics(config, samples_with_strings) - - assert "some_string" in result_2 - assert result_2["some_string"]["unique_count"] == 3 - assert result_2["some_string"]["total_count"] == 5 diff --git a/tests/unit_tests/test_dataset_viewer.py b/tests/unit_tests/test_dataset_viewer.py index b44762b4fa..cda8e3a13e 100644 --- a/tests/unit_tests/test_dataset_viewer.py +++ b/tests/unit_tests/test_dataset_viewer.py @@ -11,10 +11,18 @@ # 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 unittest.mock import mock_open, patch +from unittest.mock import patch, mock_open +from pytest import MonkeyPatch +from pydantic import BaseModel +from typing import Any -from nemo_gym.dataset_viewer import JsonlDatasetViewerConfig, build_jsonl_dataset_viewer +import json +from nemo_gym.dataset_viewer import ( + JsonlDatasetViewerConfig, + build_jsonl_dataset_viewer, + get_aggregate_metrics, +) class TestDatasetViewer: def test_sanity( @@ -31,3 +39,87 @@ def test_sanity( mock_content = r"""{"reward": 0.0, "accuracy": false, "set_overlap": 0.0, "original_term_minefield_hit": false, "order_instruction_following_failure": false, "id": 44, "expected_synonym_values": [489, 504], "expected_synonyms": ["Awake", "Alert"], "minefield_label": "Alive", "minefield_label_value": 497, "responses_create_params": {"input": [{"content": "# Instructions", "role": "system"}, {"content": "How does the human body's response to danger highlight the instinct for survival?", "role": "user"}]}, "response": {"id": "resp_689038d64ad081929f6f36d2f2554431063ce0bbdad9d001", "created_at": 1754282198.0, "error": null, "incomplete_details": null, "instructions": null, "metadata": {}, "model": "gpt-4.1-2025-04-14", "object": "response", "output": [{"content": [{"annotations": [], "text": "fake chat message", "type": "output_text"}], "role": "assistant", "type": "message", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"summary": [{"type": "summary_text", "text": "fake reasoning"}], "type": "reasoning", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"arguments": "{\"synonym\":\"Survival\"}", "call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "name": "get_synonym_value", "type": "function_call", "id": "fc_689038d5d1cc8192b91bbef0069ff82f063ce0bbdad9d001", "status": "completed"}, {"call_id": "call_pyKbpFtdag6LL6euwpAJ8UEw", "output": "{\"synonym_value\": 860}", "type": "function_call_output"}, {"arguments": "{\"synonym_values\":[860]}", "call_id": "call_N4Kr3NJJohoTaxSL5DkkG4W6", "name": "extract_synonym_values", "type": "function_call", "id": "fc_689038d6c71c8192b7ab0de7cc655d98063ce0bbdad9d001", "status": "completed"}], "parallel_tool_calls": false, "temperature": 1.0, "tool_choice": "auto", "tools": [{"name": "get_synonym_value", "parameters": {"properties": {"synonym": {"type": "string", "title": "Synonym", "description": "The synonym to get the value for."}}, "type": "object", "required": ["synonym"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Get the synonym value for a synonym.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"synonym_value\": {\"type\": \"integer\", \"title\": \"Synonym Value\", \"description\": \"The value for this synonym.\"}}, \"type\": \"object\", \"required\": [\"synonym_value\"]}\n"}, {"name": "extract_synonym_values", "parameters": {"properties": {"synonym_values": {"items": {"type": "integer"}, "type": "array", "title": "Synonym Values", "description": "The synonym values corresponding to the term for the user query."}}, "type": "object", "required": ["synonym_values"], "additionalProperties": false}, "strict": true, "type": "function", "description": "Extract the synonym values you retrieved for the term that is relevant to the user query.\nThis operation returns a value that conforms to the following JSON Schema: {\"properties\": {\"success\": {\"type\": \"boolean\", \"title\": \"Success\", \"description\": \"Success.\"}}, \"type\": \"object\", \"required\": [\"success\"]}\n"}], "top_p": 1.0, "background": false, "max_output_tokens": null, "max_tool_calls": null, "previous_response_id": null, "prompt": null, "reasoning": {"effort": null, "generate_summary": null, "summary": null}, "service_tier": "default", "status": "completed", "text": {"format": {"type": "text"}}, "top_logprobs": 0, "truncation": "disabled", "usage": {"input_tokens": 2864, "input_tokens_details": {"cached_tokens": 2798}, "output_tokens": 19, "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": 2883}, "user": null, "prompt_cache_key": null, "safety_identifier": null, "store": true}}""" with patch("builtins.open", mock_open(read_data=mock_content)): build_jsonl_dataset_viewer(config) + + def test_get_aggregate_metrics(self, monkeypatch: MonkeyPatch): + class DummySample(BaseModel): + responses_create_params: dict = {} + response: dict = {} + reward: float = 1.0 + accuracy: bool = True + set_overlap: float = 0.5 + unrelated_list: list = [] + unrelated_dict: dict = {} + + class DummySampleWithStrings(DummySample): + some_string: str + + config = JsonlDatasetViewerConfig(jsonl_fpath="") + samples = [ + DummySample(reward=1.0, accuracy=True, set_overlap=0.5), + DummySample(reward=0.0, accuracy=False, set_overlap=0.0), + DummySample(reward=0.5, accuracy=True, set_overlap=1.0), + ] + + samples_with_strings = [ + DummySampleWithStrings(reward=1.0, accuracy=True, some_string="asdf"), + DummySampleWithStrings(reward=0.0, accuracy=False, some_string="asdf"), + DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word1"), + DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word1"), + DummySampleWithStrings(reward=0.5, accuracy=True, some_string="word2"), + ] + + def mock_compute_sample_metrics(line: str): + metrics = json.loads(line) + return metrics, False + + class DummyAgg: + def model_dump(self, by_alias=True): + return {} + + class DummyDatasetMetrics: + def add(self, metrics: Any): + pass + + def aggregate(self): + return DummyAgg() + + monkeypatch.setattr( + "nemo_gym.train_data_utils.compute_sample_metrics", + mock_compute_sample_metrics, + ) + monkeypatch.setattr("nemo_gym.train_data_utils.DatasetMetrics", DummyDatasetMetrics) + + with patch("builtins.open", mock_open(read_data="{}\n")): + result_1 = get_aggregate_metrics(config, samples) + + assert "reward" in result_1 + assert "accuracy" in result_1 + assert "set_overlap" in result_1 + + assert "unrelated_str" not in result_1 + assert "unrelated_list" not in result_1 + assert "unrelated_dict" not in result_1 + + assert "responses_create_params" not in result_1 + assert "response" not in result_1 + + # Check computed values + reward_stats = result_1["reward"] + assert reward_stats["Total # non-null values"] == 3 + assert reward_stats["Average"] == (1.0 + 0.0 + 0.5) / 3 + assert reward_stats["Min"] == 0.0 + assert reward_stats["Max"] == 1.0 + + # Check computed values with bools converted to int + accuracy_stats = result_1["accuracy"] + assert accuracy_stats["Total # non-null values"] == 3 + assert accuracy_stats["Average"] == (1 + 0 + 1) / 3 + assert accuracy_stats["Min"] == 0 + assert accuracy_stats["Max"] == 1 + + with patch("builtins.open", mock_open(read_data="{}\n")): + result_2 = get_aggregate_metrics(config, samples_with_strings) + + assert "some_string" in result_2 + assert result_2["some_string"]["unique_count"] == 3 + assert result_2["some_string"]["total_count"] == 5 From 3b09058192c5212eb972d2c657ffe6c22c50c92b Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Fri, 5 Sep 2025 14:51:06 -0700 Subject: [PATCH 3/5] Format Signed-off-by: Frankie Siino --- nemo_gym/dataset_viewer.py | 46 +++----- nemo_gym/train_data_utils.py | 141 +++++++----------------- tests/unit_tests/test_dataset_viewer.py | 9 +- 3 files changed, 58 insertions(+), 138 deletions(-) diff --git a/nemo_gym/dataset_viewer.py b/nemo_gym/dataset_viewer.py index 8faa738c3f..956e3cf799 100644 --- a/nemo_gym/dataset_viewer.py +++ b/nemo_gym/dataset_viewer.py @@ -11,32 +11,27 @@ # 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 typing import List, Dict, Any - import json +from typing import Any, Dict, List -from tqdm.auto import tqdm - -from pydantic import BaseModel, ConfigDict - +from gradio import JSON, Blocks, Chatbot, ChatMessage, Dropdown +from gradio.components.chatbot import MetadataDict from openai.types.responses.response_input_param import ( - ResponseFunctionToolCallParam, - FunctionCallOutput, - ResponseReasoningItemParam, EasyInputMessageParam, + FunctionCallOutput, + ResponseFunctionToolCallParam, ResponseInputItemParam, + ResponseReasoningItemParam, ) +from pydantic import BaseModel, ConfigDict +from tqdm.auto import tqdm -from gradio import Chatbot, Blocks, ChatMessage, Dropdown, JSON -from gradio.components.chatbot import MetadataDict - -from nemo_gym.server_utils import get_global_config_dict from nemo_gym.base_resources_server import BaseVerifyResponse - +from nemo_gym.server_utils import get_global_config_dict from nemo_gym.train_data_utils import ( AvgMinMax, - compute_sample_metrics, DatasetMetrics, + compute_sample_metrics, ) @@ -89,9 +84,7 @@ def format_reasoning(m: ResponseReasoningItemParam) -> List[ChatMessage]: def format_message(m: EasyInputMessageParam) -> List[ChatMessage]: - content = ( - m["content"] if isinstance(m["content"], list) else [{"text": m["content"]}] - ) + content = m["content"] if isinstance(m["content"], list) else [{"text": m["content"]}] match m["role"]: case "user": return [ @@ -180,9 +173,7 @@ def rollout_to_messages(create_params: dict, response: dict) -> List[ChatMessage step += 1 for message in convert_single_message(m): - message.metadata["title"] = ( - f"Turn {turn} Step {step} - {message.metadata['title']}" - ) + message.metadata["title"] = f"Turn {turn} Step {step} - {message.metadata['title']}" messages.append(message) return messages @@ -248,9 +239,7 @@ def aggregate_other_metrics(data: List[DatasetViewerVerifyResponse]) -> Dict[str return result -def get_aggregate_metrics( - data: List[DatasetViewerVerifyResponse], raw_lines: List[str] -) -> Dict[str, Any]: +def get_aggregate_metrics(data: List[DatasetViewerVerifyResponse], raw_lines: List[str]) -> Dict[str, Any]: dataset_metrics = DatasetMetrics() for line in raw_lines: metrics, is_offending = compute_sample_metrics(line) @@ -271,10 +260,7 @@ def build_jsonl_dataset_viewer(config: JsonlDatasetViewerConfig) -> Blocks: raw_lines.append(line) data.append(DatasetViewerVerifyResponse.model_validate_json(line)) - choices = [ - (f"Sample {i + 1} - Responses ID {d.response.id}", i) - for i, d in enumerate(data) - ] + choices = [(f"Sample {i + 1} - Responses ID {d.response.id}", i) for i, d in enumerate(data)] def select_item(value: int): d = data[value] @@ -299,9 +285,7 @@ def select_item(value: int): layout="panel", label="Rollout", ) - item_dropdown.select( - fn=select_item, inputs=item_dropdown, outputs=chatbot, show_api=False - ) + item_dropdown.select(fn=select_item, inputs=item_dropdown, outputs=chatbot, show_api=False) return demo diff --git a/nemo_gym/train_data_utils.py b/nemo_gym/train_data_utils.py index 16536520e5..d2c40c64fa 100644 --- a/nemo_gym/train_data_utils.py +++ b/nemo_gym/train_data_utils.py @@ -11,43 +11,34 @@ # 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 typing import Dict, List, Literal, Self, Optional, Union, Tuple - +import json from abc import abstractmethod - from collections import Counter, defaultdict - -from itertools import repeat, count - -import json - -from shutil import copyfileobj - +from itertools import count, repeat from pathlib import Path +from shutil import copyfileobj +from typing import Dict, List, Literal, Optional, Self, Tuple, Union from devtools import pprint - -from tqdm.auto import tqdm - -from pydantic import BaseModel, Field, ConfigDict, ValidationError - from omegaconf import DictConfig +from pydantic import BaseModel, ConfigDict, Field, ValidationError +from tqdm.auto import tqdm +from nemo_gym.base_resources_server import BaseRunRequest from nemo_gym.config_types import ( AGENT_REF_KEY, - ServerInstanceConfig, + AgentServerRef, DatasetConfig, DatasetType, DownloadJsonlDatasetGitlabConfig, - AgentServerRef, + ServerInstanceConfig, ) +from nemo_gym.gitlab_utils import download_jsonl_dataset from nemo_gym.global_config import ( GlobalConfigDictParser, GlobalConfigDictParserConfig, get_global_config_dict, ) -from nemo_gym.gitlab_utils import download_jsonl_dataset -from nemo_gym.base_resources_server import BaseRunRequest class TrainDataProcessorConfig(BaseModel): @@ -110,19 +101,13 @@ def _aggregate(self) -> Self: class DatasetMetrics(Accumulator): number_of_examples: int = Field(serialization_alias="Number of examples", default=0) - number_of_tools: AvgMinMax = Field( - serialization_alias="Number of tools", default_factory=AvgMinMax - ) + number_of_tools: AvgMinMax = Field(serialization_alias="Number of tools", default_factory=AvgMinMax) json_dumped_number_of_words: AvgMinMax = Field( serialization_alias="Json-dumped number of words (proxy for token count)", default_factory=AvgMinMax, ) - number_of_turns: AvgMinMax = Field( - serialization_alias="Number of turns", default_factory=AvgMinMax - ) - temperature: AvgMinMax = Field( - serialization_alias="Temperature", default_factory=AvgMinMax - ) + number_of_turns: AvgMinMax = Field(serialization_alias="Number of turns", default_factory=AvgMinMax) + temperature: AvgMinMax = Field(serialization_alias="Temperature", default_factory=AvgMinMax) # TODO: Number of unique create params, Number of unique user messages, other sampling params, etc @@ -225,9 +210,7 @@ def run(self, global_config_dict: DictConfig): # pragma: no cover config = TrainDataProcessorConfig.model_validate(global_config_dict) self._print_title("Load and validate server instance configs") - server_instance_configs = self.load_and_validate_server_instance_configs( - config, global_config_dict - ) + server_instance_configs = self.load_and_validate_server_instance_configs(config, global_config_dict) self._print_title( f"Load datasets. Missing datasets {'**WILL**' if config.should_download else 'will **NOT**'} be downloaded." @@ -235,14 +218,10 @@ def run(self, global_config_dict: DictConfig): # pragma: no cover self.load_datasets(config, server_instance_configs) self._print_title("Validate samples and aggregate metrics") - dataset_type_to_aggregate_metrics = self.validate_samples_and_aggregate_metrics( - server_instance_configs - ) + dataset_type_to_aggregate_metrics = self.validate_samples_and_aggregate_metrics(server_instance_configs) self._print_title("Collate samples and aggregate metrics") - self.collate_samples( - config, server_instance_configs, dataset_type_to_aggregate_metrics - ) + self.collate_samples(config, server_instance_configs, dataset_type_to_aggregate_metrics) self._print_title("Finished!") @@ -260,19 +239,13 @@ def load_and_validate_server_instance_configs( self, config: TrainDataProcessorConfig, global_config_dict: DictConfig ) -> List[ServerInstanceConfig]: parser = GlobalConfigDictParser() - server_instance_configs = parser.filter_for_server_instance_configs( - global_config_dict - ) + server_instance_configs = parser.filter_for_server_instance_configs(global_config_dict) agent_configs: List[ServerInstanceConfig] = [ - c - for c in server_instance_configs - if c.SERVER_TYPE == "responses_api_agents" + c for c in server_instance_configs if c.SERVER_TYPE == "responses_api_agents" ] - server_names_list_str = "\n- ".join( - [""] + [f"{c.name} ({c.SERVER_TYPE})" for c in server_instance_configs] - ) + server_names_list_str = "\n- ".join([""] + [f"{c.name} ({c.SERVER_TYPE})" for c in server_instance_configs]) print( f"Found {len(server_instance_configs)} server instance configs ({len(agent_configs)} agent configs):{server_names_list_str}\n\n" ) @@ -285,9 +258,7 @@ def load_and_validate_server_instance_configs( else: agent_configs_without_data.append(agent_config) - server_names_list_str = "\n- ".join( - [""] + [f"{c.name} ({c.SERVER_TYPE})" for c in agent_configs_without_data] - ) + server_names_list_str = "\n- ".join([""] + [f"{c.name} ({c.SERVER_TYPE})" for c in agent_configs_without_data]) print( f"Found {len(agent_configs_without_data)} agent server instance configs withOUT datasets:{server_names_list_str}\n\n" ) @@ -295,9 +266,7 @@ def load_and_validate_server_instance_configs( server_names_list_str = "" for c in agent_configs_with_data: server_str = f"\n- {c.name}" - datasets_str = "\n - ".join( - [""] + [f"{d.name} ({d.type})" for d in c.datasets] - ) + datasets_str = "\n - ".join([""] + [f"{d.name} ({d.type})" for d in c.datasets]) server_names_list_str += f"{server_str}{datasets_str}" print( f"Found {len(agent_configs_with_data)} agent server instance configs WITH datasets:{server_names_list_str}\n\n" @@ -307,9 +276,7 @@ def load_and_validate_server_instance_configs( in_scope_dataset_types = config.in_scope_dataset_types agent_configs_with_in_scope_datasets: List[ServerInstanceConfig] = [] for agent_config in agent_configs_with_data: - in_scope_datasets = [ - d for d in agent_config.datasets if d.type in in_scope_dataset_types - ] + in_scope_datasets = [d for d in agent_config.datasets if d.type in in_scope_dataset_types] if not in_scope_datasets: continue @@ -320,13 +287,9 @@ def load_and_validate_server_instance_configs( server_names_list_str = "" for c in agent_configs_with_in_scope_datasets: server_str = f"\n- {c.name}" - datasets_str = "\n - ".join( - [""] + [f"{d.name} ({d.type})" for d in c.datasets] - ) + datasets_str = "\n - ".join([""] + [f"{d.name} ({d.type})" for d in c.datasets]) server_names_list_str += f"{server_str}{datasets_str}" - print( - f"In scope dataset types for `{config.mode}` mode: {in_scope_dataset_types}" - ) + print(f"In scope dataset types for `{config.mode}` mode: {in_scope_dataset_types}") print( f"Found {len(agent_configs_with_data)} agent server instance configs with in-scope datasets:{server_names_list_str}" ) @@ -351,26 +314,18 @@ def load_datasets( server_names_list_str = "" for server_name, datasets in local_datasets_found.items(): - datasets_str = "\n - ".join( - [""] + [f"{d.name} ({d.type})" for d in datasets] - ) + datasets_str = "\n - ".join([""] + [f"{d.name} ({d.type})" for d in datasets]) server_names_list_str += f"\n- {server_name}{datasets_str}" - print( - f"FOUND the following datasets at their local paths:{server_names_list_str}\n\n" - ) + print(f"FOUND the following datasets at their local paths:{server_names_list_str}\n\n") server_names_list_str = "" for server_name, datasets in local_datasets_not_found.items(): - datasets_str = "\n - ".join( - [""] + [f"{d.name} ({d.type})" for d in datasets] - ) + datasets_str = "\n - ".join([""] + [f"{d.name} ({d.type})" for d in datasets]) server_names_list_str += f"\n- {server_name}{datasets_str}" print(f"MISSING the following datasets:{server_names_list_str}\n\n") if config.mode == "example_validation": - assert not local_datasets_not_found, ( - "You must provide the above missing example jsonl files!" - ) + assert not local_datasets_not_found, "You must provide the above missing example jsonl files!" if not config.should_download: assert not local_datasets_not_found, ( "Missing local datasets. You must provide local datasets since download is disabled. Run with `+should_download=true` to enable downloading." @@ -384,9 +339,7 @@ def load_datasets( download_config = DownloadJsonlDatasetGitlabConfig.model_validate( d.gitlab_identifier.model_dump() | {"output_fpath": d.jsonl_fpath} ) - print( - f"Downloading dataset `{d.name}` from `{server_name}` using {download_config}" - ) + print(f"Downloading dataset `{d.name}` from `{server_name}` using {download_config}") download_jsonl_dataset(download_config) ######################################## @@ -422,9 +375,7 @@ def _validate_samples_and_aggregate_metrics_single_dataset( return state - def _validate_aggregate_metrics( - self, aggregate_metrics_dict: Dict, metrics_fpath: Path - ) -> Optional[Path]: + def _validate_aggregate_metrics(self, aggregate_metrics_dict: Dict, metrics_fpath: Path) -> Optional[Path]: """ Returns the conflicting metrics fpath if invalid. Else returns None """ @@ -432,9 +383,7 @@ def _validate_aggregate_metrics( with open(metrics_fpath) as f: previous_aggregate_metrics_dict = json.load(f) if aggregate_metrics_dict != previous_aggregate_metrics_dict: - conflicting_metrics_fpath = metrics_fpath.with_name( - f"{metrics_fpath.stem}_conflict.json" - ) + conflicting_metrics_fpath = metrics_fpath.with_name(f"{metrics_fpath.stem}_conflict.json") with open(conflicting_metrics_fpath, "w") as f: json.dump(aggregate_metrics_dict, f, indent=4) @@ -444,9 +393,7 @@ def validate_samples_and_aggregate_metrics( self, server_instance_configs: List[ServerInstanceConfig] ) -> Dict[str, DatasetMetrics]: conflicting_fpaths: List[str] = [] - dataset_type_to_aggregate_metrics: Dict[str, DatasetMetrics] = defaultdict( - DatasetMetrics - ) + dataset_type_to_aggregate_metrics: Dict[str, DatasetMetrics] = defaultdict(DatasetMetrics) for c in server_instance_configs: for d in c.datasets: state = self._validate_samples_and_aggregate_metrics_single_dataset(d) @@ -455,9 +402,7 @@ def validate_samples_and_aggregate_metrics( aggregate_metrics = state.metrics.aggregate() - aggregate_metrics_dict = aggregate_metrics.model_dump( - mode="json", by_alias=True - ) + aggregate_metrics_dict = aggregate_metrics.model_dump(mode="json", by_alias=True) aggregate_metrics_dict = d.model_dump() | aggregate_metrics_dict data_fpath = Path(d.jsonl_fpath) @@ -478,9 +423,7 @@ def validate_samples_and_aggregate_metrics( if conflicting_fpaths: conflicting_fpaths_str = "\n- ".join([""] + conflicting_fpaths) - raise ValueError( - f"Found conflicting aggregate metrics that need to be corrected:{conflicting_fpaths_str}" - ) + raise ValueError(f"Found conflicting aggregate metrics that need to be corrected:{conflicting_fpaths_str}") return dict(dataset_type_to_aggregate_metrics) @@ -504,9 +447,7 @@ def _collate_samples_single_type( with open(data_path) as source, open(prepare_path, "w") as target: for line in tqdm(source, desc=f"Preparing data at {data_path}"): d = json.loads(line) - d[AGENT_REF_KEY] = AgentServerRef( - type="responses_api_agents", name=c.name - ).model_dump() + d[AGENT_REF_KEY] = AgentServerRef(type="responses_api_agents", name=c.name).model_dump() target.write(f"{json.dumps(d)}\n") paths_to_collate.append(prepare_path) @@ -528,9 +469,7 @@ def collate_samples( aggregate_metrics = dataset_type_to_aggregate_metrics[type] aggregate_metrics = aggregate_metrics.aggregate() - aggregate_metrics_dict = aggregate_metrics.model_dump( - mode="json", by_alias=True - ) + aggregate_metrics_dict = aggregate_metrics.model_dump(mode="json", by_alias=True) parent = Path(config.output_dirpath) parent.mkdir(exist_ok=True) @@ -563,13 +502,9 @@ def collate_samples( if conflicting_fpaths: conflicting_fpaths_str = "\n- ".join([""] + conflicting_fpaths) - raise ValueError( - f"Found conflicting aggregate metrics that need to be corrected:{conflicting_fpaths_str}" - ) + raise ValueError(f"Found conflicting aggregate metrics that need to be corrected:{conflicting_fpaths_str}") - final_fpaths_str = "\n- ".join( - [""] + [f"{type}: {fpath}" for type, fpath in final_fpaths.items()] - ) + final_fpaths_str = "\n- ".join([""] + [f"{type}: {fpath}" for type, fpath in final_fpaths.items()]) print(f"View your final data!{final_fpaths_str}") diff --git a/tests/unit_tests/test_dataset_viewer.py b/tests/unit_tests/test_dataset_viewer.py index cda8e3a13e..d27b9ba986 100644 --- a/tests/unit_tests/test_dataset_viewer.py +++ b/tests/unit_tests/test_dataset_viewer.py @@ -11,12 +11,12 @@ # 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 unittest.mock import patch, mock_open -from pytest import MonkeyPatch -from pydantic import BaseModel +import json from typing import Any +from unittest.mock import mock_open, patch -import json +from pydantic import BaseModel +from pytest import MonkeyPatch from nemo_gym.dataset_viewer import ( JsonlDatasetViewerConfig, @@ -24,6 +24,7 @@ get_aggregate_metrics, ) + class TestDatasetViewer: def test_sanity( self, From 99afdd046804db0357d2688c63a4d1e85f49cedc Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Fri, 5 Sep 2025 15:03:51 -0700 Subject: [PATCH 4/5] Precommit fix Signed-off-by: Frankie Siino --- nemo_gym/train_data_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nemo_gym/train_data_utils.py b/nemo_gym/train_data_utils.py index d2c40c64fa..8fca05b039 100644 --- a/nemo_gym/train_data_utils.py +++ b/nemo_gym/train_data_utils.py @@ -229,9 +229,9 @@ def _print_title(self, title: str) -> None: # pragma: no cover print(f""" {"#" * 100} -# +# # {title} -# +# {"#" * 100} """) From 5f644b77cd988ea499e800a481f57b6b1f773362 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Fri, 5 Sep 2025 16:09:07 -0700 Subject: [PATCH 5/5] Fix test Signed-off-by: Frankie Siino --- tests/unit_tests/test_dataset_viewer.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/tests/unit_tests/test_dataset_viewer.py b/tests/unit_tests/test_dataset_viewer.py index d27b9ba986..403c3e682f 100644 --- a/tests/unit_tests/test_dataset_viewer.py +++ b/tests/unit_tests/test_dataset_viewer.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import json -from typing import Any from unittest.mock import mock_open, patch from pydantic import BaseModel @@ -54,7 +53,6 @@ class DummySample(BaseModel): class DummySampleWithStrings(DummySample): some_string: str - config = JsonlDatasetViewerConfig(jsonl_fpath="") samples = [ DummySample(reward=1.0, accuracy=True, set_overlap=0.5), DummySample(reward=0.0, accuracy=False, set_overlap=0.0), @@ -77,10 +75,6 @@ class DummyAgg: def model_dump(self, by_alias=True): return {} - class DummyDatasetMetrics: - def add(self, metrics: Any): - pass - def aggregate(self): return DummyAgg() @@ -88,10 +82,8 @@ def aggregate(self): "nemo_gym.train_data_utils.compute_sample_metrics", mock_compute_sample_metrics, ) - monkeypatch.setattr("nemo_gym.train_data_utils.DatasetMetrics", DummyDatasetMetrics) - with patch("builtins.open", mock_open(read_data="{}\n")): - result_1 = get_aggregate_metrics(config, samples) + result_1 = get_aggregate_metrics(samples, "{}\n") assert "reward" in result_1 assert "accuracy" in result_1 @@ -118,8 +110,8 @@ def aggregate(self): assert accuracy_stats["Min"] == 0 assert accuracy_stats["Max"] == 1 - with patch("builtins.open", mock_open(read_data="{}\n")): - result_2 = get_aggregate_metrics(config, samples_with_strings) + # Check string counts + result_2 = get_aggregate_metrics(samples_with_strings, "{}\n") assert "some_string" in result_2 assert result_2["some_string"]["unique_count"] == 3