ronantakizawa/codereview-bench huggingface dataset:
Repository: oumi-ai/oumi
File: src/oumi/datasets/sft/tulu3_sft_mixture.py (Python)

Code before:

@register_dataset("allenai/tulu-3-sft-mixture")
class Tulu3MixtureDataset(BaseSftDataset):
    default_dataset = "allenai/tulu-3-sft-mixture"

    def transform_conversation(self, example: Union[dict, pd.Series]) -> Conversation:
        """Convert the example into a Conversation.

        Args:
            example (dict or Pandas Series): An example containing a `messages`
               field which is a list of dicts with `content` and `role` string
               fields
        """
        messages = [self._transform_message(msg) for msg in example["messages"]]
        return Conversation(messages=messages)

    @classmethod
    def _transform_message(cls, message: dict) -> Message:
        roles = {"system": Role.SYSTEM, "assistant": Role.ASSISTANT, "user": Role.USER}
        content = message["content"]
        try:
            role = roles[message["role"]]
            return Message(role=role, content=content)
        except KeyError as e:
            raise ValueError(f"Unknown role {e.args[0]}, was expecting one of {
                             roles.keys()}") from e


Review comment (suggestion):
```suggestion
            raise ValueError(f"Unknown role {e.args[0]}, was expecting one of "
                             f"{roles.keys()}") from e
```

Code after:

@register_dataset("allenai/tulu-3-sft-mixture")
class Tulu3MixtureDataset(BaseSftDataset):
    default_dataset = "allenai/tulu-3-sft-mixture"

    def transform_conversation(self, example: Union[dict, pd.Series]) -> Conversation:
        """Convert the example into a Conversation.

        Args:
            example (dict or Pandas Series): An example containing a `messages`
               field which is a list of dicts with `content` and `role` string
               fields
        """
        messages = [self._transform_message(msg) for msg in example["messages"]]
        return Conversation(messages=messages)

    @classmethod
    def _transform_message(cls, message: dict) -> Message:
        roles = {"system": Role.SYSTEM, "assistant": Role.ASSISTANT, "user": Role.USER}
        content = message["content"]
        role_str = message["role"].lower().strip()
        try:
            role = roles[role_str]
            return Message(role=role, content=content)
        except KeyError as e:
            raise ValueError(
                f"Unknown role {message['role']}, "
                f"was expecting one of {roles.keys()}"
            ) from e
Repository: oumi-ai/oumi
File: tests/unit/inference/test_sambanova_inference_engine.py (Python)

Code before:

def test_get_request_headers(sambanova_engine):
    """Test generation of request headers."""
    remote_params = RemoteParams(api_key="test_api_key", api_url="<placeholder>")

    with patch.object(
        SambanovaInferenceEngine,
        "_get_api_key",
        return_value="test_api_key",
    ):
        result = sambanova_engine._get_request_headers(remote_params)

    assert result["Content-Type"] == "application/json"
    assert result["X-API-Key"] == "test_api_key"


def test_get_supported_params(sambanova_engine):
    """Test supported generation parameters."""
    supported_params = sambanova_engine.get_supported_params()

    assert supported_params == {
        "max_new_tokens",
        "stop_strings",
        "temperature",
        "top_p",
    }


Review comment (suggestion):
These parameters should match those in the `get_supported_params` method in sambanova_inference_engine.py

Please update either the test or the engine to reflect the parameters users can set.

Code after:

def test_get_request_headers(sambanova_engine):
    """Test generation of request headers."""
    remote_params = RemoteParams(api_key="test_api_key", api_url="<placeholder>")

    with patch.object(
        SambanovaInferenceEngine,
        "_get_api_key",
        return_value="test_api_key",
    ):
        result = sambanova_engine._get_request_headers(remote_params)

    assert result["Content-Type"] == "application/json"
    assert result["Authorization"] == "Bearer test_api_key"


def test_get_supported_params(sambanova_engine):
    """Test supported generation parameters."""
    supported_params = sambanova_engine.get_supported_params()

    assert supported_params == {
        "max_new_tokens",
        "stop_strings",
        "temperature",
        "top_p",
    }
Repository: oumi-ai/oumi
File: src/oumi/cli/fetch.py (Python)

Code before:

    else:
        # raise error
        logger.error("Invalid config path")
        raise typer.Exit(1)

    try:
        # Fetch from GitHub
        github_url = f"{OUMI_GITHUB_RAW}/{config_path}"
        response = requests.get(github_url)
        response.raise_for_status()
        config_content = response.text

        # Validate YAML
        yaml.safe_load(config_content)

        # Save to destination
        local_path = (config_dir or Path(OUMI_DIR).expanduser()) / config_path
        local_path.parent.mkdir(parents=True, exist_ok=True)

        with open(local_path, "w") as f:
            f.write(config_content)

        logger.info(f"Successfully downloaded config to {local_path}")

    except requests.RequestException as e:
        logger.error(f"Failed to download config from GitHub: {e}")
        raise typer.Exit(1)
    except yaml.YAMLError:
        logger.error("Invalid YAML configuration")
        raise typer.Exit(1)


Review comment (bug):
One edge case to consider: what if the user already fetched the file?

I would recommend calling `local_path.exists()` and if it does, log a warning that we're downloading a config that already exists.

Code after:
    """Fetch configuration files from GitHub repository."""
    # Remove oumi:// prefix if present
    config_path, config_dir = resolve_oumi_prefix(config_path, output_dir)

    try:
        # Check destination first
        local_path = (config_dir or Path(OUMI_DIR).expanduser()) / config_path
        if local_path.exists() and not force:
            msg = f"Config already exists at {local_path}. Use --force to overwrite"
            logger.error(msg)
            typer.echo(msg, err=True)
            raise typer.Exit(code=1)

        # Fetch from GitHub
        github_url = f"{OUMI_GITHUB_RAW}/{config_path.lstrip('/')}"
        response = requests.get(github_url)
        response.raise_for_status()
        config_content = response.text

        # Validate YAML
        yaml.safe_load(config_content)

        # Save to destination
        if local_path.exists():
            logger.warning(f"Overwriting existing config at {local_path}")
        local_path.parent.mkdir(parents=True, exist_ok=True)

        with open(local_path, "w") as f:
            f.write(config_content)
        logger.info(f"Successfully downloaded config to {local_path}")

    except RequestException as e:
        logger.error(f"Failed to download config from GitHub: {e}")
        raise typer.Exit(1)
    except yaml.YAMLError:
        logger.error("Invalid YAML configuration")
        raise typer.Exit(1)
Repository: oumi-ai/oumi
File: src/oumi/cli/cli_utils.py (Python)

Code before:

    Returns:
        tuple[str, Path]: (cleaned path, output directory)
    """
    config_path = config_path[7:]

    config_dir = output_dir or os.environ.get("OUMI_DIR") or "~/.oumi/configs"
    config_dir = Path(config_dir).expanduser()
    config_dir.mkdir(parents=True, exist_ok=True)

    return config_path, config_dir


def resolve_and_fetch_config(
    config_path: str, output_dir: Optional[Path] = None
) -> Path:
    """Resolve oumi:// prefix and fetch config if needed.

    Args:
        config_path: Original config path that may contain oumi:// prefix
        output_dir: Optional override for output directory

    Returns:
        Path: Local path to the config file
    """
    if not config_path.startswith("oumi://"):
        return Path(config_path)

    from oumi.cli.fetch import fetch

    fetch(config_path, output_dir)

    return Path(config_path)


Review comment (suggestion):
```suggestion
    if not config_path.lower().startswith("oumi://"):
```

Code after:

    Returns:
        tuple[str, Path]: (cleaned path, output directory)
    """
    oumi_prefix = "oumi://"
    if config_path.lower().startswith(oumi_prefix):
        config_path = config_path[len(oumi_prefix) :]

    config_dir = output_dir or os.environ.get("OUMI_DIR") or "~/.oumi/configs"
    config_dir = Path(config_dir).expanduser()
    config_dir.mkdir(parents=True, exist_ok=True)

    return config_path, config_dir


def resolve_and_fetch_config(
    config_path: str, output_dir: Optional[Path] = None
) -> Path:
    """Resolve oumi:// prefix and fetch config if needed.

    Args:
        config_path: Original config path that may contain oumi:// prefix
        output_dir: Optional override for output directory

    Returns:
        Path: Local path to the config file
    """
    if not config_path.lower().startswith("oumi://"):
        return Path(config_path)

    from oumi.cli.fetch import fetch

    fetch(config_path, output_dir)

    return Path(config_path)
Repository: oumi-ai/oumi
File: src/oumi/cli/main.py (Python)

Code before:
    launch_app.command(help="Turns down a cluster.")(down)
    launch_app.command(
        name="run", context_settings=CONTEXT_ALLOW_EXTRA_ARGS, help="Runs a job."
    )(launcher_run)
    launch_app.command(help="Prints the status of jobs launched from Oumi.")(status)
    launch_app.command(help="Stops a cluster.")(stop)
    launch_app.command(
        context_settings=CONTEXT_ALLOW_EXTRA_ARGS, help="Launches a job."
    )(up)
    launch_app.command(help="Prints the available clouds.")(which)
    app.add_typer(launch_app, name="launch", help="Launch jobs remotely.")

    distributed_app = typer.Typer(pretty_exceptions_enable=False)
    distributed_app.command(context_settings=CONTEXT_ALLOW_EXTRA_ARGS)(accelerate)
    distributed_app.command(context_settings=CONTEXT_ALLOW_EXTRA_ARGS)(torchrun)
    app.add_typer(
        distributed_app,
        name="distributed",
        help=(
            "A wrapper for torchrun/accelerate "
            "with reasonable default values for distributed training."
        ),
    )

    app.command(
        help="Fetch configuration files from GitHub repository.",
    )(fetch)

    return app


def run():
    """The entrypoint for the CLI."""
    app = get_app()
    return app()


if "sphinx" in sys.modules:
    # Create the CLI app when building the docs to auto-generate the CLI reference.
    app = get_app()


Review comment (suggestion):
```suggestion
        help="Fetch configuration files from the oumi GitHub repository.",
```

Small nit :)

Code after:
    launch_app.command(help="Turns down a cluster.")(down)
    launch_app.command(
        name="run", context_settings=CONTEXT_ALLOW_EXTRA_ARGS, help="Runs a job."
    )(launcher_run)
    launch_app.command(help="Prints the status of jobs launched from Oumi.")(status)
    launch_app.command(help="Stops a cluster.")(stop)
    launch_app.command(
        context_settings=CONTEXT_ALLOW_EXTRA_ARGS, help="Launches a job."
    )(up)
    launch_app.command(help="Prints the available clouds.")(which)
    app.add_typer(launch_app, name="launch", help="Launch jobs remotely.")

    distributed_app = typer.Typer(pretty_exceptions_enable=False)
    distributed_app.command(context_settings=CONTEXT_ALLOW_EXTRA_ARGS)(accelerate)
    distributed_app.command(context_settings=CONTEXT_ALLOW_EXTRA_ARGS)(torchrun)
    app.add_typer(
        distributed_app,
        name="distributed",
        help=(
            "A wrapper for torchrun/accelerate "
            "with reasonable default values for distributed training."
        ),
    )

    app.command(
        help="Fetch configuration files from the oumi GitHub repository.",
    )(fetch)

    return app


def run():
    """The entrypoint for the CLI."""
    app = get_app()
    return app()


if "sphinx" in sys.modules:
    # Create the CLI app when building the docs to auto-generate the CLI reference.
    app = get_app()
Repository: oumi-ai/oumi
File: src/oumi/datasets/vision_language/pixmo_cap.py (Python)

Code before:
    Type,
)


@register_dataset("allenai/pixmo-cap")
class PixmoCapDataset(VisionLanguageSftDataset):
    """Dataset class for the `allenai/pixmo-cap` dataset."""

    default_dataset = "allenai/pixmo-cap"

    @override
    def transform_conversation(self, example: dict) -> Conversation:
        """Transform the example into a Conversation object."""
        messages: list[Message] = []
        messages.append(
            Message(
                role=Role.USER,
                content=[
                    ContentItem(type=Type.IMAGE_URL, content=example["image_url"])
                ],
            )
        )
        messages.append(Message(role=Role.ASSISTANT, content=example["caption"]))
        for transcript in example["transcripts"]:
            messages.append(Message(role=Role.ASSISTANT, content=transcript))

        return Conversation(messages=messages)


Review comment (suggestion):
is there a USER text query ?   

If not, we may want to add a hardcoded taskspeciffc prompt,  something like this:  https://github.com/oumi-ai/oumi/blob/3720d477900026fd47e8ce7e2562c0e5f981e648/src/oumi/datasets/vision_language/flickr30k.py#L37 
to give LLM some instruction what to do .

Code after:
    Type,
)


@register_dataset("allenai/pixmo-cap")
class PixmoCapDataset(VisionLanguageSftDataset):
    """Dataset class for the `allenai/pixmo-cap` dataset.

    The dataset is affected by some image URLs having a 404 issue.
    """

    default_dataset = "allenai/pixmo-cap"

    @override
    def transform_conversation(self, example: dict) -> Conversation:
        """Transform the example into a Conversation object.

        A "transcripts" column is also available but not used yet.
        """
        input_text = "Describe this image:"

        messages: list[Message] = []
        messages.append(
            Message(
                role=Role.USER,
                content=[
                    ContentItem(type=Type.IMAGE_URL, content=example["image_url"]),
                    ContentItem(type=Type.TEXT, content=input_text),
                ],
            )
        )
        messages.append(Message(role=Role.ASSISTANT, content=example["caption"]))

        return Conversation(messages=messages)
Repository: oumi-ai/oumi
File: src/oumi/core/configs/training_config.py (Python)

Code before:
                self.data.train.datasets,
                self.data.validation.datasets,
                self.data.test.datasets,
            ):
                if (
                    "processor_name" not in dataset_params.dataset_kwargs
                    or "processor_kwargs" in dataset_params.dataset_kwargs
                ):
                    continue
                dataset_processor_name: str = dataset_params.dataset_kwargs[
                    "processor_name"
                ]
                if dataset_processor_name == model_processor_name:
                    # Copy processor kwargs from the model if processor names match
                    # and the dataset doesn't override them.
                    dataset_params.dataset_kwargs["processor_kwargs"] = {
                        **self.model.processor_kwargs
                    }

        # Verl will error without a validation dataset.
        if (
            self.training.trainer_type == TrainerType.VERL_GRPO
            and not self.data.validation.datasets
        ):
            raise ValueError(
                "At least one validation dataset is required for VERL GRPO training."
            )


Review comment (suggestion):
```suggestion
                "At least one validation dataset is required for VERL_GRPO training."
```
Minor nit for consistency with other messages

Code after:
                self.data.train.datasets,
                self.data.validation.datasets,
                self.data.test.datasets,
            ):
                if (
                    "processor_name" not in dataset_params.dataset_kwargs
                    or "processor_kwargs" in dataset_params.dataset_kwargs
                ):
                    continue
                dataset_processor_name: str = dataset_params.dataset_kwargs[
                    "processor_name"
                ]
                if dataset_processor_name == model_processor_name:
                    # Copy processor kwargs from the model if processor names match
                    # and the dataset doesn't override them.
                    dataset_params.dataset_kwargs["processor_kwargs"] = {
                        **self.model.processor_kwargs
                    }

        # Verl will error without a validation dataset.
        if (
            self.training.trainer_type == TrainerType.VERL_GRPO
            and not self.data.validation.datasets
        ):
            raise ValueError(
                "At least one validation dataset is required for VERL_GRPO training."
            )
Repository: oumi-ai/oumi
File: src/oumi/train.py (Python)

Code before:
# Copyright 2025 - Oumi
#
# 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 functools
import time
from importlib.metadata import version
from pathlib import Path
from pprint import pformat
from typing import Any, Callable, Final, Optional, Union

try:
    import ray  # pyright: ignore[reportMissingImports]
except ModuleNotFoundError:
    ray = None
import torch
import transformers
from transformers.trainer_utils import get_last_checkpoint

from oumi.builders import (
    build_collator_from_config,
    build_dataset_mixture,
    build_metrics_function,
    build_model,
    build_peft_model,
    build_processor,
    build_reward_functions,
    build_tokenizer,
    build_trainer,
    build_training_callbacks,
    is_image_text_llm,
)
from oumi.core.configs import (
    DatasetSplit,
    TrainerType,
    TrainingConfig,
)
from oumi.core.configs.internal.supported_models import (
    is_custom_model,
)

Review comment (suggestion):
Instead of importing `ray` at the top of `train.py`, I'd suggest importing it inside the body of `_verl_train`. That way it's clear that this is a `verl` dependency, not an `oumi` dependency.

Code after:
# Copyright 2025 - Oumi
#
# 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 functools
import time
from importlib.metadata import version
from pathlib import Path
from pprint import pformat
from typing import Any, Callable, Final, Optional, Union

import torch
import transformers
from transformers.trainer_utils import get_last_checkpoint

from oumi.builders import (
    build_collator_from_config,
    build_dataset_mixture,
    build_metrics_function,
    build_model,
    build_peft_model,
    build_processor,
    build_reward_functions,
    build_tokenizer,
    build_trainer,
    build_training_callbacks,
    is_image_text_llm,
)
from oumi.core.configs import (
    DatasetSplit,
    TrainerType,
    TrainingConfig,
)
from oumi.core.configs.internal.supported_models import (
    is_custom_model,
)
from oumi.core.distributed import (
    barrier,
    cleanup_distributed,
    estimate_dataloader_num_workers,
Repository: oumi-ai/oumi
File: tests/e2e/test_train_e2e.py (Python)

Code before:
            ),
            batch_size=16,
            dataloader_num_workers=2,
            dataloader_prefetch_factor=4,
            max_steps=20,
        ),
        TrainTestConfig(
            test_name="train_text_smollm_135m_sft",
            config_path=(
                get_configs_dir() / "recipes" / "smollm" / "sft" / "135m" / "train.yaml"
            ),
            max_steps=10,
        ),
    ],
    ids=get_train_test_id_fn,
)
@pytest.mark.e2e
@pytest.mark.single_gpu
def test_train_text_1gpu_24gb(
    test_config: TrainTestConfig,
    tmp_path: Path,
):
    _test_train_impl(test_config=test_config, tmp_path=tmp_path, use_distributed=False)


@requires_gpus(count=4, min_gb=40.0)
@pytest.mark.parametrize(
    "test_config",
    [
        TrainTestConfig(
            test_name="train_mm_qwen2_vl_2b_trl_sft_fft",
            config_path=(
                get_configs_dir()
                / "recipes"
                / "vision"
                / "qwen2_vl_2b"
                / "sft"
                / "full"
                / "train.yaml"
            ),
            trainer_type=TrainerType.TRL_SFT,
            max_steps=5,
            save_steps=5,
        ),
        TrainTestConfig(
            test_name="train_mm_qwen2_vl_2b_oumi_fft",
            config_path=(
                get_configs_dir()
                / "recipes"
                / "vision"
                / "qwen2_vl_2b"

Review comment (suggestion):
I previously had to use value 39 (instead of 40) because not 100% of memory is available (and/or  1024 vs 1000 difference? don't remember exactly) 

Code after:
            ),
            batch_size=16,
            dataloader_num_workers=2,
            dataloader_prefetch_factor=4,
            max_steps=20,
        ),
        TrainTestConfig(
            test_name="train_text_smollm_135m_sft",
            config_path=(
                get_configs_dir() / "recipes" / "smollm" / "sft" / "135m" / "train.yaml"
            ),
            max_steps=10,
        ),
    ],
    ids=get_train_test_id_fn,
)
@pytest.mark.e2e
@pytest.mark.single_gpu
def test_train_text_1gpu_24gb(
    test_config: TrainTestConfig,
    tmp_path: Path,
):
    _test_train_impl(test_config=test_config, tmp_path=tmp_path, use_distributed=False)


@requires_gpus(count=4, min_gb=39.0)
@pytest.mark.parametrize(
    "test_config",
    [
        TrainTestConfig(
            test_name="train_mm_qwen2_vl_2b_trl_sft_fft",
            config_path=(
                get_configs_dir()
                / "recipes"
                / "vision"
                / "qwen2_vl_2b"
                / "sft"
                / "full"
                / "train.yaml"
            ),
            trainer_type=TrainerType.TRL_SFT,
            max_steps=5,
            save_steps=5,
        ),
        TrainTestConfig(
            test_name="train_mm_qwen2_vl_2b_oumi_fft",
            config_path=(
                get_configs_dir()
                / "recipes"
                / "vision"
                / "qwen2_vl_2b"