From 618d57f23c1c63c354d0886c03bbb62f085e2cd9 Mon Sep 17 00:00:00 2001
From: Colin <colin>
Date: Tue, 26 Mar 2024 18:15:55 +0800
Subject: [PATCH] Update define.

---
 wit/meaning_dataset.py | 9 ++++++++-
 wit/train.py           | 9 +++++----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/wit/meaning_dataset.py b/wit/meaning_dataset.py
index 1e89819..40d8b87 100644
--- a/wit/meaning_dataset.py
+++ b/wit/meaning_dataset.py
@@ -31,6 +31,7 @@ class MeaningMap:  # 16777216 1048576 8192
             self.ms_data = np.load(file_data)
             self.ms_start = np.load(file_start)
             self.ms_len = np.load(file_len)
+            print("Load end")
         else:
             print("Disk cache miss, build new one.")
 
@@ -123,6 +124,11 @@ class MeaningDataset(Dataset):
                 self.data.append(sq)
                 self.length.append(len(sq))
 
+        unique, counts = np.unique(self.length, return_counts=True)
+        print("MeaningDataset size: " + str(len(self.length)))
+        print("MeaningDataset max sequence length: " + str(max(unique)))
+        print("MeaningDataset most popular sequence length: " + str(unique[np.argmax(counts)]))
+
     def __len__(self):
         return len(self.data)
 
@@ -205,7 +211,8 @@ if __name__ == "__main__":
     md = MeaningDataset(4096, 8100, size=1024)
     train, val = md.Split(0.95)
 
-    dl = BatchGroupMeaningDataloader(train, 2)
+    dl = BatchGroupMeaningDataloader(train, 32)
+    length = len(dl)
     it = iter(dl)
     ne1 = next(it)
     ne2 = next(it)
diff --git a/wit/train.py b/wit/train.py
index 54eb5f6..79d958b 100644
--- a/wit/train.py
+++ b/wit/train.py
@@ -18,15 +18,16 @@ learning_rate = 0.0001
 use_tril_attention_mask = None
 precision = "32-true"  # "precision:bf16-mixed,16-mixed,32-true"
 train_batch_size = 32
-val_batch_size = 32
+val_batch_size = 4
 num_proc = 8
 max_epochs = 1000
 strategy = "auto"
 resume_from_ckpt_path = None
 seed = 42
 
-vocab_size = 1024
+vocab_size = 2048
 level_ratio = 4
+level = 4
 
 hidden_size = 256  # 128 1024 2048  32
 num_attention_heads = 8  # 8 8 16
@@ -57,9 +58,9 @@ if __name__ == "__main__":
     lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
     tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
 
-    start = vocab_size * level_ratio * level_ratio * level_ratio * level_ratio
+    start = vocab_size * (level_ratio**level)
     end = start * level_ratio
-    size = start + start
+    size = vocab_size * (level_ratio ** (level / 2))
     raw_dataset = MeaningDataset(start, end, size, vocab_size, level_ratio)
     train_dataset, val_dataset = raw_dataset.Split(0.95)
     train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)