fix: 修复预处理流水线 4 个致命问题
- pipeline: 单条数据异常不再崩溃整条流水线, 改 log warning 后跳过 - pipeline: _align_bucket 统一用 len(ids) 填充, 修复多输出模式下长度错配 - writer: BinWriter/H5Writer 写入失败自动清理残留文件并记录详细错误 - packing: BFDPacking 真正将序列打包进 bin 而非仅重排, 减少碎片
This commit is contained in:
@@ -6,8 +6,7 @@ pipeline later flattens the result into contiguous tensors.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Tuple
|
||||
from typing import Dict, List
|
||||
|
||||
from astrai.factory import BaseFactory
|
||||
|
||||
@@ -53,6 +52,15 @@ class SimplePacking(PackingStrategy):
|
||||
|
||||
@PackingStrategyFactory.register("bfd")
|
||||
class BFDPacking(PackingStrategy):
|
||||
"""Best-Fit Decreasing bin packing.
|
||||
|
||||
Assigns sequences to bins using a best-fit heuristic (sorted by
|
||||
decreasing length) and concatenates sequences within each bin into
|
||||
a single packed sequence. Packed sequences are truncated to
|
||||
*max_packed_len* so that each packed bin fits within one context
|
||||
window during training.
|
||||
"""
|
||||
|
||||
def apply(
|
||||
self,
|
||||
keys: Dict[str, List[List[int]]],
|
||||
@@ -62,24 +70,40 @@ class BFDPacking(PackingStrategy):
|
||||
sequences = keys.get("sequence", [])
|
||||
if not sequences:
|
||||
return keys
|
||||
plan = self._plan(sequences, max_packed_len)
|
||||
reordered: dict = defaultdict(list)
|
||||
for orig_idx, _ in plan:
|
||||
for k, vals in keys.items():
|
||||
reordered[k].append(
|
||||
_truncate(vals[orig_idx], max_packed_len, truncation_mode)
|
||||
bins = self._plan(sequences, max_packed_len, truncation_mode)
|
||||
|
||||
packed: Dict[str, List[List[int]]] = {}
|
||||
for k, vals in keys.items():
|
||||
packed[k] = [
|
||||
_truncate(
|
||||
self._concat_bin(vals, bin_indices),
|
||||
max_packed_len,
|
||||
truncation_mode,
|
||||
)
|
||||
return dict(reordered)
|
||||
for bin_indices in bins
|
||||
]
|
||||
return packed
|
||||
|
||||
@staticmethod
|
||||
def _plan(sequences: List[List[int]], max_packed_len: int) -> List[Tuple[int, int]]:
|
||||
def _concat_bin(vals: List[List[int]], indices: List[int]) -> List[int]:
|
||||
result: List[int] = []
|
||||
for i in indices:
|
||||
result.extend(vals[i])
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _plan(
|
||||
sequences: List[List[int]], max_packed_len: int, truncation_mode: str
|
||||
) -> List[List[int]]:
|
||||
n = len(sequences)
|
||||
order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
|
||||
bins: List[List[int]] = []
|
||||
bin_lengths: List[int] = []
|
||||
|
||||
for orig_idx in order:
|
||||
seq_len = min(len(sequences[orig_idx]), max_packed_len)
|
||||
seq_len = len(
|
||||
_truncate(sequences[orig_idx], max_packed_len, truncation_mode)
|
||||
)
|
||||
best_bin = None
|
||||
best_remain = max_packed_len + 1
|
||||
for i, bl in enumerate(bin_lengths):
|
||||
@@ -94,8 +118,4 @@ class BFDPacking(PackingStrategy):
|
||||
bins.append([orig_idx])
|
||||
bin_lengths.append(seq_len)
|
||||
|
||||
plan: List[Tuple[int, int]] = []
|
||||
for bin_indices in bins:
|
||||
for orig_idx in bin_indices:
|
||||
plan.append((orig_idx, min(len(sequences[orig_idx]), max_packed_len)))
|
||||
return plan
|
||||
return bins
|
||||
|
||||
Reference in New Issue
Block a user