ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(utils): 优化序列打包逻辑以提高效率和正确性

Browse Source
This commit is contained in:
ViperEkura
2025-10-07 12:11:40 +08:00
parent cd3a22c1ac
commit 491b39b23f
1 changed files with 14 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/utils.py
+10 -23
View File
@@ -39,37 +39,24 @@ def comprehensive_normalization(text):
def pack_sequences(sequences: List[Tensor], pack_size: int, pad_value: int) -> List[Tensor]: def pack_sequences(sequences: List[Tensor], pack_size: int, pad_value: int) -> List[Tensor]:
packages = [] packages = []
sequences.sort(key=lambda x: x.numel(), reverse=True) sequences.sort(key=lambda x: x.numel(), reverse=True)
current_pack = torch.tensor([], dtype=torch.int32) current_pack = torch.full((pack_size,), pad_value, dtype=torch.int32)
current_pos = 0
for tensor in sequences: for tensor in sequences:
if tensor.numel() > pack_size: if tensor.numel() > pack_size:
packages.append(tensor[:pack_size]) tensor = tensor[:pack_size]
continue
remaining = pack_size - current_pack.numel() tensor_size = tensor.numel()
if remaining == 0: if current_pos + tensor_size > pack_size:
packages.append(current_pack) packages.append(current_pack)
current_pack = tensor current_pack = torch.full((pack_size,), pad_value, dtype=torch.int32)
elif tensor.numel() <= remaining: current_pos = 0
current_pack = torch.cat([current_pack, tensor])
else:
padding = torch.full((remaining,), pad_value, dtype=torch.int32)
current_pack = torch.cat([current_pack, padding])
packages.append(current_pack)
current_pack = tensor
if current_pack.numel() > 0: current_pack[current_pos: current_pos + tensor_size] = tensor
if current_pack.numel() < pack_size: current_pos += tensor_size
padding = torch.full(
(pack_size - current_pack.numel(),),
pad_value,
dtype=torch.int32
)
current_pack = torch.cat([current_pack, padding])
else:
current_pack = current_pack[:pack_size]
if current_pos > 0:
packages.append(current_pack) packages.append(current_pack)
return packages return packages