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

pipeline/io/export.py

@@ -12,7 +12,7 @@ from tqdm import tqdm
 from pipeline.io.file_scanner import FileScanner
 from pipeline.io.hdf5_handler import HDF5Handler
 from pipeline.processors import BaseProcessor
-from pipeline.packing import SequencePacker
+from pipeline.packing import pack_tensors
 from pipeline.utils import error_handler
 
 logger = logging.getLogger(__name__)
@@ -127,10 +127,12 @@ def cache_jsonl(
                     continue
 
         if pack_size > 0:
-            output = {}
+            dtypes = (
-            for key in output_keys:
+                dict(processor.schema.output_fields)
-                packer = SequencePacker(pack_size, pad_value)
+                if processor.schema is not None
-                output[key] = packer.pack(arrows[key])
+                else None
+            )
+            output = pack_tensors(arrows, pack_size, pad_value, dtypes)
         else:
             output = arrows
 

pipeline/packing.py

@@ -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