guidellm.data.preprocessors.mappers

@PreprocessorRegistry.register("generative_column_mapper")
class GenerativeColumnMapper(DataDependentPreprocessor):
    defaults: ClassVar[dict[str, list[str]]] = {
        "prompt_tokens_count_column": ["prompt_tokens_count", "input_tokens_count"],
        "output_tokens_count_column": [
            "output_tokens_count",
            "completion_tokens_count",
        ],
        "prefix_column": [
            "system_prompt",
            "system",
            "prefix",
        ],
        "text_column": [
            "prompt",
            "instruction",
            "question",
            "input",
            "context",
            "content",
            "conversation",
            "turn",
            "text",
        ],
        "image_column": [
            "image",
            "picture",
            "photo",
            "img",
        ],
        "video_column": [
            "video",
            "clip",
            "movie",
            "footage",
            "mp4",
            "mov",
            "avi",
        ],
        "audio_column": [
            "audio",
            "sound",
            "voice",
            "speech",
            "wav",
            "mp3",
        ],
        "tools_column": [
            "tools",
            "functions",
            "tool_definitions",
        ],
        "tool_response_column": [
            "tool_response",
            "tool_result",
            "tool_output",
        ],
        "relative_timestamp_column": ["relative_timestamp"],
        "requeue_delay_column": ["requeue_delay"],
    }
    column_name_pattern: str = (
        r"^(?P<full_name>(?P<match_name>({name})(es|s)?)([-_](?P<turn>\d+))?)$"
    )

    @staticmethod
    def _filter_for_dataset(names: list[str], *dataset_names: str) -> list[str]:
        filtered_names: list[str] = []
        for name in names:
            if "." in name:
                dataset_part, column_part = name.split(".", 1)
                if dataset_part in dataset_names:
                    filtered_names.append(column_part)
            else:
                filtered_names.append(name)

        return filtered_names

    @staticmethod
    def _extract_turn_columns(
        turn_pattern: str, columns_str: str
    ) -> list[tuple[int, str]]:
        # Now find all columns that match a variant of the base name
        turn_matches = re.finditer(turn_pattern, columns_str, re.M | re.I)

        turn_columns: list[tuple[int, str]] = []
        turn_count = 0
        for match in turn_matches:
            column_name = match.group("full_name")
            if not column_name:
                continue

            turn_str = match.group("turn")
            turn = int(turn_str) if turn_str is not None else turn_count
            turn_columns.append((turn, column_name))
            turn_count += 1

        return turn_columns

    @classmethod
    def datasets_mappings(
        cls,
        datasets: list[Dataset | IterableDataset],
        input_mappings: dict[str, str | list[str]] | None = None,
    ) -> dict[DatasetColumnKey, list[DatasetColumnValue]]:
        """
        Resolve column mappings across one or more datasets.

        For each dataset, matches actual column names against the requested
        mapping names (or :attr:`defaults`) using regex patterns that account
        for pluralisation and turn suffixes (e.g. ``prompt-0``, ``prompt-1``).

        :param datasets: The loaded datasets to inspect for column names.
        :param input_mappings: Optional explicit column mappings. When ``None``,
            :attr:`defaults` is used. Values may be a single name or a list of
            candidate names in priority order.
        :return: A dict keyed by ``(column_type, turn_index)`` whose values are
            lists of ``(dataset_index, column_name)`` pairs indicating where
            each logical column can be found. Categories with no matching
            columns are silently omitted from the result.
        """
        mappings: dict[DatasetColumnKey, list[DatasetColumnValue]] = defaultdict(list)
        input_map: dict[str, list[str]] = cls.defaults
        if input_mappings:
            input_map = {
                k: v if isinstance(v, list) else [v] for k, v in input_mappings.items()
            }

        for index, dataset in enumerate(datasets):
            dataset_name = (
                dataset.info.dataset_name
                if dataset.info and dataset.info.dataset_name
                else index
            )
            dataset_columns = dataset.column_names or list(next(iter(dataset)).keys())
            dataset_columns_str = "\n".join(dataset_columns)

            for column_type, names in input_map.items():
                filtered_names = cls._filter_for_dataset(
                    names, str(index), str(dataset_name)
                )
                if not filtered_names:
                    continue

                column_pattern = cls.column_name_pattern.format(
                    name="|".join(re.escape(n) for n in filtered_names)
                )
                # Find the first matching column name
                base_match = re.search(column_pattern, dataset_columns_str, re.M | re.I)
                if not base_match:
                    continue

                turn_pattern = cls.column_name_pattern.format(
                    name=base_match.group("match_name"),
                )
                turn_columns = cls._extract_turn_columns(
                    turn_pattern,
                    dataset_columns_str,
                )

                for turn, column_name in sorted(turn_columns):
                    column_type = cast("GenerativeDatasetColumnType", column_type)
                    mappings[(column_type, turn)].append((index, column_name))

        return mappings

    def __init__(
        self,
        config: GenerativeColumnMapperArgs,
    ):
        self.input_mappings = config.column_mappings
        self.datasets_column_mappings: (
            dict[DatasetColumnKey, list[DatasetColumnValue]] | None
        )

    def __call__(self, items: list[dict[str, Any]]) -> list[dict[str, list[Any]]]:
        if self.datasets_column_mappings is None:
            raise ValueError("DefaultGenerativeColumnMapper not setup with data.")

        mapped: list[dict[str, Any]] = []

        for (column_type, turn), column_mappings in sorted(
            self.datasets_column_mappings.items()
        ):
            # Ensure the mapped list has enough turns for this turn
            # Should never need to happen
            while len(mapped) <= turn:
                mapped.append(defaultdict(list))

            for (
                dataset_index,
                dataset_column,
            ) in column_mappings:
                mapped[turn][column_type].append(
                    items[dataset_index]["dataset"][dataset_column]
                )

        return [dict(m) for m in mapped if len(m) > 0]

    def setup_data(
        self,
        datasets: list[DatasetType],
    ):
        self.datasets_column_mappings = self.datasets_mappings(
            datasets, self.input_mappings
        )

        if not self.datasets_column_mappings:
            raise ValueError(
                "GenerativeColumnMapper found no matching columns. "
                f"Requested mappings: {self.input_mappings or 'default mappings'}. "
                "Every row will produce an empty result."
            )

@DataPreprocessorArgs.register(
    [
        "generative_column_mapper",
        "pooling_column_mapper",
    ]
)
class GenerativeColumnMapperArgs(DataPreprocessorArgs):
    """Model for generative column mapper preprocessor arguments."""

    kind: Literal["generative_column_mapper", "pooling_column_mapper"] = Field(
        default="generative_column_mapper",
        description="Type identifier for the generative column mapper preprocessor.",
    )
    column_mappings: dict[str, str | list[str]] | None = Field(
        default=None,
        description="Mappings for the column names.",
        examples=[
            {
                "prompt_tokens_count_column": [
                    "prompt_tokens_count",
                    "input_tokens_count",
                ],
                "output_tokens_count_column": [
                    "output_tokens_count",
                    "completion_tokens_count",
                ],
            }
        ],
    )