From 7a8815cceb09918d87b5043fd4ca709160f5f736 Mon Sep 17 00:00:00 2001
From: Colin <colin>
Date: Fri, 29 Mar 2024 22:10:25 +0800
Subject: [PATCH] Refine the base code.

---
 wit/lit_module.py      |  3 ++-
 wit/meaning_dataset.py | 11 ++++++++---
 wit/train.py           | 30 +++++++++++-------------------
 3 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/wit/lit_module.py b/wit/lit_module.py
index ced90cb..af2fc4d 100644
--- a/wit/lit_module.py
+++ b/wit/lit_module.py
@@ -93,5 +93,6 @@ class LitModule(pl.LightningModule):
             stopping_threshold=1,
         )
         lr_monitor = pl.callbacks.LearningRateMonitor(logging_interval="step")
-        return [checkpoint_callback, lr_monitor]
+        return [lr_monitor]
+        # return [checkpoint_callback, lr_monitor]
         # return [checkpoint_callback, early_stop_callback]
diff --git a/wit/meaning_dataset.py b/wit/meaning_dataset.py
index 40d8b87..9d52122 100644
--- a/wit/meaning_dataset.py
+++ b/wit/meaning_dataset.py
@@ -125,9 +125,12 @@ class MeaningDataset(Dataset):
                 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)]))
+        print("----------------------------------------------------------------")
+        print("MeaningDataset start:" + str(start) + " end:" + str(end) + " space:" + str(end - start))
+        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)]))
+        print("----------------------------------------------------------------")
 
     def __len__(self):
         return len(self.data)
@@ -197,6 +200,8 @@ class BatchGroupMeaningDataloader(Dataset):
             np.random.shuffle(index_shuffle)
             index = index[index_shuffle]
         self.indexBatch = index
+        print("Dataloader batch size:" + str(batch_size) + " count:" + str(len(index)))
+        print("Dataloader total:" + str(len(length)) + " drop:" + str(len(length) - len(index) * batch_size))
 
     def __len__(self):
         return len(self.indexBatch)
diff --git a/wit/train.py b/wit/train.py
index 79d958b..955c13d 100644
--- a/wit/train.py
+++ b/wit/train.py
@@ -17,34 +17,26 @@ pretrain_model_name = None  # "qwen/Qwen-1_8B-Chat"
 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 = 4
+train_batch_size = 4
+val_batch_size = 1
 num_proc = 8
 max_epochs = 1000
 strategy = "auto"
 resume_from_ckpt_path = None
 seed = 42
 
-vocab_size = 2048
+vocab_size = 1024
 level_ratio = 4
 level = 4
+dataset_level = 1
 
 hidden_size = 256  # 128 1024 2048  32
 num_attention_heads = 8  # 8 8 16
-num_hidden_layers = 1  # 6 12 24  3
+num_hidden_layers = 2  # 6 12 24  3
 
-name = "vocab_level_hidden_head_layer"
-version = (
-    str(vocab_size)
-    + "_"
-    + str(level_ratio)
-    + "_"
-    + str(hidden_size)
-    + "_"
-    + str(num_attention_heads)
-    + "_"
-    + str(num_hidden_layers)
-)
+name = "vocab_ratio_level_data_hidden_head_layer"
+ver = f"{vocab_size}" + "_" + f"{level_ratio}" + "_" + f"{level}" + "_" + f"{dataset_level}"
+ver = ver + "_" + f"{hidden_size}" + "_" + f"{num_attention_heads}" + "_" + f"{num_hidden_layers}"
 
 if __name__ == "__main__":
     torch.manual_seed(seed)
@@ -60,9 +52,9 @@ if __name__ == "__main__":
 
     start = vocab_size * (level_ratio**level)
     end = start * level_ratio
-    size = vocab_size * (level_ratio ** (level / 2))
+    size = int(vocab_size * (level_ratio**dataset_level))
     raw_dataset = MeaningDataset(start, end, size, vocab_size, level_ratio)
-    train_dataset, val_dataset = raw_dataset.Split(0.95)
+    train_dataset, val_dataset = raw_dataset.Split(0.9)
     train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
     val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)
     # it = iter(train_dataloader)
@@ -75,7 +67,7 @@ if __name__ == "__main__":
         accelerator="cuda",
         devices=[0, 1],
         precision=precision,
-        logger=TBLogger("./log/", name=name, version=version, default_hp_metric=False),
+        logger=TBLogger("./log/", name=name, version=ver, default_hp_metric=False),
         strategy=strategy,
         max_epochs=max_epochs,
     )