33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
from typing import List
|
|
|
|
import torch
|
|
|
|
|
|
def tensorize_batch(sequences: List[torch.Tensor], padding_value, align_right=False) -> torch.Tensor:
|
|
if len(sequences[0].size()) == 1:
|
|
max_len_1 = max([s.size(0) for s in sequences])
|
|
out_dims = (len(sequences), max_len_1)
|
|
out_tensor = sequences[0].new_full(out_dims, padding_value)
|
|
for i, tensor in enumerate(sequences):
|
|
length_1 = tensor.size(0)
|
|
if align_right:
|
|
out_tensor[i, -length_1:] = tensor
|
|
else:
|
|
out_tensor[i, :length_1] = tensor
|
|
return out_tensor
|
|
elif len(sequences[0].size()) == 2:
|
|
max_len_1 = max([s.size(0) for s in sequences])
|
|
max_len_2 = max([s.size(1) for s in sequences])
|
|
out_dims = (len(sequences), max_len_1, max_len_2)
|
|
out_tensor = sequences[0].new_full(out_dims, padding_value)
|
|
for i, tensor in enumerate(sequences):
|
|
length_1 = tensor.size(0)
|
|
length_2 = tensor.size(1)
|
|
if align_right:
|
|
out_tensor[i, -length_1:, :length_2] = tensor
|
|
else:
|
|
out_tensor[i, :length_1, :length_2] = tensor
|
|
return out_tensor
|
|
else:
|
|
raise
|