fix: mask 全链路保持 bool dtype(创建→打包→HDF5落盘/读盘)

This commit is contained in:
2026-05-15 17:23:51 +08:00
parent 625695fd71
commit e1125be3a7
2 changed files with 54 additions and 8 deletions
+47 -3
View File
@@ -1,7 +1,9 @@
import logging
from typing import List
from typing import Any, Dict, List, Optional, Tuple, Union
import torch
from torch import Tensor
from pipeline.utils import error_handler
logger = logging.getLogger(__name__)
@@ -38,12 +40,15 @@ class SequencePacker:
"""
def __init__(
self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = torch.int32
self,
pack_size: int,
pad_value: Union[int, bool] = 0,
dtype: Optional[torch.dtype] = None,
):
self.pack_size = pack_size
self.pad_value = pad_value
self.dtype = dtype
self._buffer: List[int] = []
self._buffer: List = []
self._pos: int = 0
self._packages: List[Tensor] = []
@@ -61,6 +66,8 @@ class SequencePacker:
Sequences are concatenated in order and sliced at pack_size boundaries.
The final chunk is padded with pad_value.
When dtype is not set at init, it is inferred from the first input tensor.
Args:
sequences: List of 1D input tensors.
@@ -70,6 +77,10 @@ class SequencePacker:
if not sequences:
return []
# --- auto-infer dtype from first sequence ---
if self.dtype is None:
self.dtype = sequences[0].dtype
# --- validate & normalize ---
normalized: List[Tensor] = []
for i, seq in enumerate(sequences):
@@ -100,3 +111,36 @@ class SequencePacker:
self._pos = len(buf)
return self._packages
def pack_tensors(
tensors: Dict[str, List[Tensor]],
pack_size: int,
pad_value: Union[int, bool] = 0,
dtypes: Optional[Dict[str, torch.dtype]] = None,
) -> Dict[str, List[Tensor]]:
"""
Pack multiple named tensor groups in parallel.
Each group is packed independently with its own SequencePacker instance.
When dtypes is provided, packers use the declared dtype per key;
otherwise dtype is auto-inferred from the first tensor in each group.
Args:
tensors: Dict mapping key names to lists of 1D tensors.
pack_size: Fixed chunk length.
pad_value: Padding value for non-bool tensors.
dtypes: Optional per-key dtype declarations.
Returns:
Dict mapping key names to lists of packed tensors.
"""
if dtypes is None:
dtypes = {}
output: Dict[str, List[Tensor]] = {}
for key, seqs in tensors.items():
dtype = dtypes.get(key)
packer = SequencePacker(pack_size, pad_value, dtype=dtype)
output[key] = packer.pack(seqs)
return output