Diff Coverage

        value = item.get(key)
        if value is not None:
            out[key] = torch.as_tensor(value, dtype=dtype)
    if "num_items_in_batch" in item:
        out["num_items_in_batch"] = int(item["num_items_in_batch"])
    return out


class TokenizedDataset(Dataset):

                    break
                yielded += 1
                yield _tensorize_item(item)
            if not made_progress:
                break

    def __len__(self) -> int:
        if self.logical_length is None:
            raise TypeError(
                "Streaming dataset has no static length. Set data.train_size "
                "to define a logical per-epoch sample budget."
            )
        return self.logical_length

    if data_type == "json_file":
        return load_dataset("json", data_files=train_path, split="train", streaming=streaming)
    subset = args.data.subset
    if subset:
        return load_dataset(train_path, subset, split="train", streaming=streaming)
    return load_dataset(train_path, split="train", streaming=streaming)


def _maybe_truncate(ds: Any, args: Any) -> Any:

def _maybe_clip_max_steps(*, base: Any, args: Any, global_sample_count: Optional[int]) -> None:
    """Clip ``base.state.max_steps`` when a finite global dataset size is known."""
    if global_sample_count is None:
        return
    steps_per_epoch = global_sample_count // max(args.train.global_batch_size, 1)
    num_epochs = max(int(args.train.num_train_epochs or 1), 1)
    base.state.max_steps = min(args.train.max_steps, num_epochs * steps_per_epoch)

def _per_rank_train_size(train_size: Optional[int], dp_rank: int, dp_size: int) -> Optional[int]:
    """Split a global logical sample budget across DP ranks."""
    if train_size is None:
        return None
    if train_size < 0:
        raise ValueError(f"data.train_size must be >= 0, but got {train_size}")
    base = train_size // max(dp_size, 1)
    remainder = train_size % max(dp_size, 1)
    return base + (1 if dp_rank < remainder else 0)

            "desc": "Tokenizing",
        }
        try:
            ds = ds.map(data_transform, **map_kwargs)
        except TypeError as exc:
            if "unexpected keyword argument 'desc'" not in str(exc):
                raise
            map_kwargs.pop("desc", None)
            ds = ds.map(data_transform, **map_kwargs)
    column_names = getattr(ds, "column_names", []) or []
    if "input_ids" in column_names:
        ds = ds.filter(lambda x: len(x["input_ids"]) > 0)
    if args.data.shuffle:

            "Streaming dataset logical epoch budget: global=%d, local=%d",
            global_train_size, local_train_size,
        )
    else:
        logger.info(
            "Streaming dataset uses source exhaustion for epoch boundaries; "
            "set data.train_size to define a logical sample budget."
        )
    return StreamingTokenizedDataset(ds, logical_length=local_train_size)

        """
        if getattr(self, "train_dataset", None) is not None:
            return
        data_type = self.args.data.type
        dp_rank, dp_size = self._get_dp_shard_rank_and_size()
        self.train_dataset = build_dataset(
            data_type,
            base=self,
            args=self.args,

        micro_bs = self.args.train.micro_batch_size

        # Sampler uses DP rank/size — TP/CP/PP/EP peers share data.
        dp_rank, dp_size = self._get_dp_shard_rank_and_size()

        # StatefulDataLoader supports state_dict() / load_state_dict()
        # for checkpoint resume (torchdata API, used by  + ).
        num_workers = self.args.data.num_workers

                num_workers,
            )
            num_workers = 0

        is_iterable_dataset = isinstance(self.train_dataset, IterableDataset)
        loader_kwargs = {
            "batch_size": micro_bs,
            "collate_fn": self.collate_fn,
            "num_workers": num_workers,

            "num_workers": num_workers,
            "pin_memory": pin_memory,
            "drop_last": True,
        }
        if not is_iterable_dataset:
            shuffle = self.args.data.shuffle
            sampler_seed = self.args.train.seed
            self.sampler = DistributedSampler(
                self.train_dataset,
                num_replicas=dp_size,
                rank=dp_rank,
                shuffle=shuffle,

                shuffle=shuffle,
                seed=sampler_seed,
                drop_last=True,
            )
            loader_kwargs["sampler"] = self.sampler
        elif hasattr(self, 'sampler'):
            delattr(self, 'sampler')
        # prefetch_factor is only accepted when num_workers > 0
        if num_workers > 0 and prefetch_factor is not None:
            loader_kwargs["prefetch_factor"] = prefetch_factor
        if self._deterministic:

            self.args.train.global_batch_size // (micro_bs * dp_size),
            1,
        )

        try:
            dataset_size = len(self.train_dataset)
        except TypeError:
            dataset_size = None
        if dataset_size is None:
            logger.info_rank0(
                "Dataloader built: micro_bs=%d, grad_accum=%d, dataset_size=streaming/unknown",
                micro_bs, self._grad_accum,
            )
        else:
            logger.info_rank0(
                "Dataloader built: micro_bs=%d, grad_accum=%d, dataset_size=%d",
                micro_bs, self._grad_accum, dataset_size,
            )

        for epoch in range(num_epochs):
            if self.state.global_step >= self.state.max_steps:
                break
            self.state.epoch = epoch
            if hasattr(self.train_dataset, 'set_epoch'):
                self.train_dataset.set_epoch(epoch)
            if hasattr(self, 'sampler'):
                self.sampler.set_epoch(epoch)
            self.on_epoch_begin()

            return None

    def _get_dp_shard_rank_and_size(self) -> tuple[int, int]:
        """Return the DP rank/size used to partition training data."""
        dp_size = self.parallel_dims.dp_size
        non_dp = self.parallel_dims.non_dp_size
        global_rank = platform.get_rank()
        try:
            dp_rank = self.mesh["dp"].get_local_rank()
        except (KeyError, ValueError, RuntimeError):
            dp_rank = global_rank // non_dp if non_dp > 1 else global_rank
        return dp_rank, dp_size

    def _build_fsdp_kwargs(self) -> dict:
        """Build kwargs for ``fully_shard`` calls (dense parameters).

    """Normalize one-or-many media references to a list."""
    if value is None:
        return []
    if isinstance(value, (list, tuple)):
        return list(value)
    return [value]


def _iter_batched_rows(examples: Dict[str, List[Any]]) -> List[Dict[str, Any]]:

def _iter_batched_rows(examples: Dict[str, List[Any]]) -> List[Dict[str, Any]]:
    """Convert an HF ``map(batched=True)`` payload into row dictionaries."""
    if not examples:
        return []
    first_key = next(iter(examples))
    size = len(examples[first_key])
    return [{key: value[idx] for key, value in examples.items()} for idx in range(size)]

    video_queue = _as_media_list(row.get("videos", row.get("video")))

    def _normalize_content(content: Any) -> List[Dict[str, Any]]:
        if isinstance(content, str):
            return [{"type": "text", "text": content}]
        if not isinstance(content, list):
            raise ValueError(
                "VL HuggingFace rows must provide list-style message content "
                f"or plain text, but got {type(content)}."
            )
        out: List[Dict[str, Any]] = []

            )
        out: List[Dict[str, Any]] = []
        for item in content:
            if isinstance(item, str):
                out.append({"type": "text", "text": item})
                continue
            if not isinstance(item, dict):
                raise ValueError(
                    "VL message content items must be dict/str, "
                    f"but got {type(item)}."
                )
            item_type = item.get("type", "text")

            if item_type == "image":
                image = item.get("image", item.get("path", item.get("url")))
                if image is None:
                    if not image_queue:
                        raise ValueError(
                            "VL message declared an image placeholder but no "
                            f"row-level media was found in '{image_key}'."
                        )
                    image = image_queue.pop(0)

                        )
                    image = image_queue.pop(0)
                out.append({"type": "image", "image": image})
                continue
            if item_type == "video":
                video = item.get("video", item.get("path", item.get("url")))
                if video is None:
                    if not video_queue:
                        raise ValueError(
                            "VL message declared a video placeholder but no "
                            "row-level video field was found."
                        )
                    video = video_queue.pop(0)
                out.append({"type": "video", "video": video})
                continue
            raise ValueError(f"Unsupported VL content type: {item_type}")
        return out

    if row.get(messages_key) is not None:
        messages = row[messages_key]

    if row.get(messages_key) is not None:
        messages = row[messages_key]
        if not isinstance(messages, list):
            raise ValueError(
                f"data.messages_key='{messages_key}' must point to a list, "
                f"but got {type(messages)}."
            )
        return [

    user_content: List[Dict[str, Any]] = []
    for image in image_queue:
        user_content.append({"type": "image", "image": image})
    for video in video_queue:
        user_content.append({"type": "video", "video": video})
    if user_text is not None:
        user_content.append({"type": "text", "text": str(user_text)})
    if not user_content:
        raise ValueError(
            "VL HuggingFace row must provide either messages, text/instruction, "
            "or media fields."
        )

    def _squeeze_sequence(value: Any) -> Any:
        if torch.is_tensor(value):
            if value.dim() == 2 and value.shape[0] == 1:
                return value[0]
            if value.dim() == 3 and value.shape[1] == 1:
                return value[:, 0]
        return value

    def _squeeze_media(value: Any) -> Any:
        if torch.is_tensor(value) and value.dim() >= 3 and value.shape[0] == 1:
            return value[0]
        return value

    def _to_python(value: Any) -> Any:
        if torch.is_tensor(value):

    def _to_python(value: Any) -> Any:
        if torch.is_tensor(value):
            return value.detach().cpu().tolist()
        if hasattr(value, "tolist") and not isinstance(value, (list, tuple, dict, str)):
            return value.tolist()
        return value

    input_ids = _squeeze_sequence(processed["input_ids"])
    attention_mask = _squeeze_sequence(processed.get("attention_mask"))
    input_ids_list = _to_python(input_ids)

            continue
        if key in ("pixel_values", "pixel_values_videos", "image_grid_thw", "video_grid_thw"):
            value = _squeeze_media(value)
        else:
            value = _squeeze_sequence(value)
        record[key] = _to_python(value)
    return record

                return_tensors="pt",
            )
            record = _normalize_vl_processor_output(processed)
            if max_seq_len and len(record["input_ids"]) > max_seq_len:
                raise ValueError(
                    "VL HuggingFace sample length exceeds data.max_seq_len. "
                    "Automatic truncation is disabled because it can desync "
                    "vision placeholders and media features."
                )

        self.base.tokenizer = getattr(self.base.processor, "tokenizer", None)
        logger.info("Processor loaded from %s", processor_path)

    def _build_data_transform(self):
        if self.base.processor is None:
            self.base.data_transform = None
            return
        self.base.data_transform = _build_vl_hf_transform(
            self.base.processor, self.base.args.data,
        )

    @staticmethod