File size: 2,942 Bytes
ad16788 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import logging
from typing import Iterator
from typing import Tuple
from typeguard import check_argument_types
from espnet2.fileio.read_text import read_2column_text
from espnet2.samplers.abs_sampler import AbsSampler
class UnsortedBatchSampler(AbsSampler):
"""BatchSampler with constant batch-size.
Any sorting is not done in this class,
so no length information is required,
This class is convenient for decoding mode,
or not seq2seq learning e.g. classification.
Args:
batch_size:
key_file:
"""
def __init__(
self,
batch_size: int,
key_file: str,
drop_last: bool = False,
utt2category_file: str = None,
):
assert check_argument_types()
assert batch_size > 0
self.batch_size = batch_size
self.key_file = key_file
self.drop_last = drop_last
# utt2shape:
# uttA <anything is o.k>
# uttB <anything is o.k>
utt2any = read_2column_text(key_file)
if len(utt2any) == 0:
logging.warning(f"{key_file} is empty")
# In this case the, the first column in only used
keys = list(utt2any)
if len(keys) == 0:
raise RuntimeError(f"0 lines found: {key_file}")
category2utt = {}
if utt2category_file is not None:
utt2category = read_2column_text(utt2category_file)
if set(utt2category) != set(keys):
raise RuntimeError(
f"keys are mismatched between {utt2category_file} != {key_file}"
)
for k, v in utt2category.items():
category2utt.setdefault(v, []).append(k)
else:
category2utt["default_category"] = keys
self.batch_list = []
for d, v in category2utt.items():
category_keys = v
# Apply max(, 1) to avoid 0-batches
N = max(len(category_keys) // batch_size, 1)
if not self.drop_last:
# Split keys evenly as possible as. Note that If N != 1,
# the these batches always have size of batch_size at minimum.
cur_batch_list = [
category_keys[i * len(keys) // N : (i + 1) * len(keys) // N]
for i in range(N)
]
else:
cur_batch_list = [
tuple(category_keys[i * batch_size : (i + 1) * batch_size])
for i in range(N)
]
self.batch_list.extend(cur_batch_list)
def __repr__(self):
return (
f"{self.__class__.__name__}("
f"N-batch={len(self)}, "
f"batch_size={self.batch_size}, "
f"key_file={self.key_file}, "
)
def __len__(self):
return len(self.batch_list)
def __iter__(self) -> Iterator[Tuple[str, ...]]:
return iter(self.batch_list)
|