Repository: oumi-ai/oumi
File: src/oumi/core/trainers/verl_grpo_trainer.py (Python)

Code before:
    ) -> dict:
        prompt, images, answer = (
            VerlGrpoTrainer._extract_question_images_answer_from_single_turn_conversation(
                example
            )
        )
        data = {
            "data_source": data_source,
            "prompt": [
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            "images": images,
            "ability": "math",
            "reward_model": {"style": "rule", "ground_truth": answer},
            "extra_info": {
                "split": split,
                "index": idx,
                "answer": answer,
                "question": prompt,  # TODO: extract problem
            },
        }
        return data

    def _create_dataset_files(
        self, process_fn: Optional[Callable[[dict, int, str, str], dict]] = None
    ) -> None:
        """Creates dataset files for verl in Parquet format.

        The Parquet files are saved to the Oumi cache directory.
        """
        train_file = self._cache_dir / "train.parquet"
        train_dataset = self._train_dataset
        # Limit the max number of sub-processes
        num_proc = min(8, os.cpu_count() or 1)

        if process_fn is not None:
            train_data_source = self._get_data_source_name(self._oumi_config.data.train)
            train_dataset = train_dataset.map(
                function=lambda example, idx: process_fn(
                    example,
                    idx,
                    train_data_source,
                    "train",
                ),
                with_indices=True,
                num_proc=num_proc,
            )


Review comment (suggestion):
In general, I'm worried that the dataset preprocessing logic being added may be too hardcoded, and won't generalize well to the different types of datasets people may want to use here. This logic mainly seems to be for VLM examples which include a single prompt, answer, and image, right? Could we instead have this logic live in the registered geometry dataset, similar to what I did for Countdown? If VLM GRPO isn't supported by TRL yet, I think there isn't an issue with the processing logic being hardcoded for verl.

Code after:
        if len(user_messages) != 1:
            raise ValueError(f"Expected 1 user message, but got {len(user_messages)}.")

        assistant_messages = conversation.filter_messages(
            role=ConversationRole.ASSISTANT
        )
        if len(assistant_messages) != 1:
            raise ValueError(
                f"Expected 1 assistant message, but got {len(assistant_messages)}."
            )

        user_message = user_messages[0]
        assistant_message = assistant_messages[0]
        prompt: str = user_message.text_content_items[-1].content or ""
        images = [{"bytes": item.binary} for item in user_message.image_content_items]
        answer: str = assistant_message.text_content_items[-1].content or ""

        if len(images) > 0:
            # TODO: Generalize. This only works for QwenVL 2.5, which is the only
            # VLM supported by verl as of 2025-05-15.
            if not prompt.startswith("<image>"):
                prompt = "<image>" + prompt
        return (prompt, images, answer)

    @staticmethod
    def _create_verl_data_entry_from_single_turn_conversation(
        example: dict, idx: int, data_source: str, split: str
    ) -> dict:
        prompt, images, answer = (
            VerlGrpoTrainer._extract_question_images_answer_from_single_turn_conversation(
                example
            )
        )
        data = {
            "data_source": data_source,
            "prompt": [
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            "images": images,
            "ability": "math",
            "reward_model": {"style": "rule", "ground_truth": answer},
            "extra_info": {
                "split": split,
                "index": idx,
                "answer": answer,
                "question": prompt,  # TODO: extract problem
            },
        }