Add auto size support in meaning dataset.

This commit is contained in:
Colin 2024-04-23 19:05:06 +08:00
parent 24957aa2ae
commit 50e502ae96
1 changed files with 5 additions and 2 deletions

View File

@ -256,8 +256,8 @@ class MeaningDataset(Dataset):
self, self,
start, start,
end, end,
size,
vocab_size, vocab_size,
size=None,
max_subitem=10, max_subitem=10,
min_subitem=1, min_subitem=1,
min_seq_len=2, min_seq_len=2,
@ -278,7 +278,10 @@ class MeaningDataset(Dataset):
self.seq_meaning = [] self.seq_meaning = []
self.m_height = map.ms_height self.m_height = map.ms_height
self.m_weight = map.ms_weight self.m_weight = map.ms_weight
if size:
meanings = np.random.randint(start, end, size=(size)) meanings = np.random.randint(start, end, size=(size))
else:
meanings = np.arange(start, end)
seq_len = [] seq_len = []
for m in meanings: for m in meanings: