From 50e502ae96d2d5542404f3dec2f1c3c4ba681f33 Mon Sep 17 00:00:00 2001
From: Colin <colin>
Date: Tue, 23 Apr 2024 19:05:06 +0800
Subject: [PATCH] Add auto size support in meaning dataset.

---
 wit/meaning_dataset.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/wit/meaning_dataset.py b/wit/meaning_dataset.py
index 5def565..a702250 100644
--- a/wit/meaning_dataset.py
+++ b/wit/meaning_dataset.py
@@ -256,8 +256,8 @@ class MeaningDataset(Dataset):
         self,
         start,
         end,
-        size,
         vocab_size,
+        size=None,
         max_subitem=10,
         min_subitem=1,
         min_seq_len=2,
@@ -278,7 +278,10 @@ class MeaningDataset(Dataset):
         self.seq_meaning = []
         self.m_height = map.ms_height
         self.m_weight = map.ms_weight
-        meanings = np.random.randint(start, end, size=(size))
+        if size:
+            meanings = np.random.randint(start, end, size=(size))
+        else:
+            meanings = np.arange(start, end)
 
         seq_len = []
         for m in meanings: