[tab] miles/utils/data_new.py (python)
import json
import logging
import os
import random
import re

import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import itertools
import ray

from miles.utils.types import MultimodalTypes, Sample

from .timer import Timer

__all__ = ["Dataset"]

logger = logging.getLogger(__name__)


# TODO: don't read the whole file into memory.
def read_file(path):
    path, row_slice = _parse_generalized_path(path)

    if not os.path.exists(path):
        raise FileNotFoundError(f"Prompt dataset path '{path}' does not exist.")

    if path.endswith(".jsonl"):
        yield from _read_jsonl(path, row_slice)
    elif path.endswith(".parquet"):
        yield from _read_parquet(path, row_slice)
    else:
        raise ValueError(f"Unsupported file format: {path}. Supported formats are .jsonl and .parquet.")


def _read_jsonl(path, row_slice):
    # Count lines if necessary for negative slicing or logging
    total_rows = None
    if row_slice is not None:
        with open(path, "rb") as f:
            total_rows = sum(1 for _ in f)
        logger.info(f"read_file path={path} slice {total_rows=} rows into {row_slice=}")
        start, stop, step = row_slice.indices(total_rows)
    else:
        start, stop, step = 0, None, 1

    with open(path, "r", encoding="utf-8") as f:
        iterator = itertools.islice(f, start, stop, step)
        for line in iterator:
            try:
                data = json.loads(line)
            except json.JSONDecodeError:
                continue
            
            # emulate pd.read_json(dtype={"label": str})
            if "label" in data:
                data["label"] = str(data["label"])
            yield data


def _read_parquet(path, row_slice):
    pf = pq.ParquetFile(path)
    total_rows = pf.metadata.num_rows

    if row_slice is not None:
        logger.info(f"read_file path={path} slice {total_rows=} rows into {row_slice=}")
        start, stop, step = row_slice.indices(total_rows)
    else:
        start, stop, step = 0, total_rows, 1

    current_idx = 0
    # Step is always 1 based on _parse_generalized_path
    
    for batch in pf.iter_batches():
        batch_len = batch.num_rows
        batch_end = current_idx + batch_len

        # Skip batch if completely before start
        if batch_end <= start:
            current_idx = batch_end
            continue
        
        # Stop if completely after stop (and step is positive default)
        if current_idx >= stop:
            break

        rows = batch.to_pylist()
        
        # Calculate start index within this batch
        local_start = max(0, start - current_idx)
        # Calculate end index within this batch
        local_end = min(batch_len, stop - current_idx)
        
        for i in range(local_start, local_end):
             # Handle step if it was possible (it's 1 here so trivial)
             # But if general slicing was needed:
             # global_idx = current_idx + i
             # if (global_idx - start) % step != 0: continue
             
             row = rows[i]
             if "label" in row and row["label"] is not None:
                 row["label"] = str(row["label"])
             yield row

        current_idx = batch_end


def _parse_generalized_path(s: str):
    if (m := re.match(r"^(?P<real_path>.*)@\[(?P<start>-?\d*):(?P<end>-?\d*)\]$", s)) is not None:
        path = m.group("real_path")
        start = int(x) if (x := m.group("start")) != "" else None
        end = int(x) if (x := m.group("end")) != "" else None
        return path, slice(start, end)

    return s, None


def _should_skip_prompt(prompt, tokenizer, processor, max_length, apply_chat_template_kwargs):
    if max_length is None:
        return False

    from miles.utils.processing_utils import prepare_model_inputs

    input_ids, _ = prepare_model_inputs(prompt, tokenizer, processor, None, apply_chat_template_kwargs)
    return len(input_ids) > max_length


def _build_messages(data: dict, prompt_key: str, multimodal_keys: dict = None):
    messages = data.get(prompt_key)

    if isinstance(messages, str):
        messages = [{"role": "user", "content": messages}]

    if multimodal_keys:
        # Build mapping: placeholder -> (MultimodalType, content_list)
        multimodals = {}
        for type_name, data_key in multimodal_keys.items():
            mt = MultimodalTypes.get(type_name)
            if mt:
                multimodals[mt.placeholder] = (mt, list(data.get(data_key)))

        pattern = "(" + "|".join(re.escape(p) for p in multimodals.keys()) + ")"

        for message in messages:
            if isinstance(message["content"], str):
                content_list = []
                for segment in re.split(pattern, message["content"]):
                    if not segment:
                        continue
                    if segment in multimodals:
                        mt, content = multimodals[segment]
                        content_list.append({"type": mt.name, mt.name: content.pop(0)})
                    else:
                        content_list.append({"type": "text", "text": segment})
                message["content"] = content_list

            elif isinstance(message["content"], list):
                # TODO: handle more general cases. where message['content'] is a dict and contains multiple types of content.
                # e.g.
                #  "content": [
                #     {
                #         "type": "image",
                #         "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
                #     },
                #     {"type": "text", "text": "Describe this image."},
                # ],
                logger.warning("message['content'] is a list of dicts, no processing will be done.")
                continue
            else:
                raise ValueError(
                    f"Unsupported content type: {type(message['content'])}, expected str or list of dicts"
                )

    return messages


class Dataset:
    def __init__(
        self,
        path,
        tokenizer,
        processor,
        max_length,
        *,
        prompt_key="text",
        multimodal_keys=None,
        label_key=None,
        tool_key=None,
        metadata_key="metadata",
        seed=42,
        apply_chat_template=False,
        apply_chat_template_kwargs=None,
    ):
        self.origin_samples = []
        for data in read_file(path):
            prompt = _build_messages(data, prompt_key, multimodal_keys)

            metadata = data.get(metadata_key) or {}
            if tool_key is not None and tool_key in data:
                tools = data[tool_key]
                if isinstance(tools, str):
                    tools = json.loads(tools)
                elif isinstance(tools, np.ndarray):
                    tools = tools.tolist()
                assert isinstance(tools, list), f"tools must be a list, got {type(tools)} instead"
                metadata["tools"] = tools

            # TODO: this is slow.
            if _should_skip_prompt(prompt, tokenizer, processor, max_length, apply_chat_template_kwargs):
                continue

            self.origin_samples.append(
                Sample(
                    prompt=prompt,
                    label=data[label_key] if label_key is not None else None,
                    metadata=metadata,
                )
            )

        self.epoch_id = -1
        self.seed = seed
        self.samples = self.origin_samples

    def shuffle(self, new_epoch_id):
        if self.epoch_id == new_epoch_id:
            return

        random.seed(self.seed + new_epoch_id)
        permutation = list(range(len(self.samples)))
        random.shuffle(permutation)
        self.samples = [self.origin_samples[i] for i in permutation]
        self.epoch_id = new_epoch_id

    def __getitem__(self, idx):
        return self.samples[idx]

    def __len__(self):
        return len(self.samples)


def get_minimum_num_micro_batch_size(total_lengths, max_tokens_per_gpu):
    # use first fit to get the number of micro batches
    batches = []
    for length in total_lengths:
        for i in range(len(batches)):
            if batches[i] + length <= max_tokens_per_gpu:
                batches[i] += length
                break
        else:
            batches.append(length)

    return len(batches)


def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size):
    assert len(rollout_data_ref) == dp_size
    rollout_data = ray.get(rollout_data_ref[dp_rank].inner)

    partition = rollout_data.pop("partition")
    total_lengths = rollout_data["total_lengths"]

    # save the seqlen of the whole rollout batch
    Timer().seq_lens = total_lengths
    rollout_data["total_lengths"] = [total_lengths[i] for i in partition]

    return rollout_data