use custom vocab_size.

.gitignore

@@ -3,6 +3,6 @@ __pycache__
 *.txt
 *.npy
 temp
-# lightning_logs
+lightning_logs
 
 checkpoints

wit/configuration.py

@@ -4,12 +4,12 @@
 # LICENSE file in the root directory of this source tree.
 
 
-class QWenConfig:
+class ModelConfig:
     def __init__(self):
         self.vocab_size = 4096
-        self.hidden_size = 128  # 128 1024 2048  32
-        self.num_hidden_layers = 6  # 6 12 24  3
-        self.num_attention_heads = 8  # 8 8 16
+        self.hidden_size = 1024
+        self.num_hidden_layers = 24
+        self.num_attention_heads = 16
         self.emb_dropout_prob = 0.0
         self.attn_dropout_prob = 0.0
         self.layer_norm_epsilon = 1e-6

wit/demo.py

@@ -4,7 +4,7 @@ from modelscope import snapshot_download
 
 from modeling_wit import QWenLMHeadModel
 from modeling_wit import QwenRunner
-from configuration_qwen import QWenConfig
+from wit.configuration import ModelConfig
 from tokenization_qwen import QWenTokenizer
 
 
@@ -20,7 +20,7 @@ torch.cuda.manual_seed_all(seed)
 model_dir = snapshot_download("qwen/Qwen-1_8B-Chat")
 # model_dir = "/home/colin/.cache/modelscope/hub/qwen/Qwen-1_8B-Chat"
 
-config = QWenConfig()
+config = ModelConfig()
 model = QWenLMHeadModel(config)
 
 print(model)

wit/lightning_logs/version_0/hparams.yaml

@@ -1,3 +0,0 @@
-model_dir: /home/colin/.cache/modelscope/hub/qwen/Qwen-1_8B-Chat
-learning_rate: 0.0001
-use_tril_attention_mask: null

wit/lightning_logs/version_1/hparams.yaml

@@ -1,3 +0,0 @@
-model_dir: /home/colin/.cache/modelscope/hub/qwen/Qwen-1_8B-Chat
-learning_rate: 0.0001
-use_tril_attention_mask: null

wit/lightning_logs/version_2/hparams.yaml

@@ -1,3 +0,0 @@
-model_dir: /home/colin/.cache/modelscope/hub/qwen/Qwen-1_8B-Chat
-learning_rate: 0.0001
-use_tril_attention_mask: null

wit/lit_module.py

@@ -6,7 +6,7 @@ import torch
 import torchmetrics
 
 from modeling_wit import QWenLMHeadModel
-from configuration_qwen import QWenConfig
+from wit.configuration import ModelConfig
 
 from transformers import AutoConfig
 
@@ -14,15 +14,20 @@ from transformers import AutoConfig
 class LitModule(pl.LightningModule):
     def __init__(
         self,
-        model_dir: str,
+        pretrained_model_dir: str = None,
         learning_rate: float = 0.0001,
+        config: ModelConfig = None,
         use_tril_attention_mask: str = False,
     ):
         super().__init__()
         self.save_hyperparameters()
-        config = QWenConfig()
+        if config == None:
+            config = ModelConfig()
         model = QWenLMHeadModel(config)
-        model = model.from_pretrained(model_dir)
+        if pretrained_model_dir != None:
+            from modelscope import snapshot_download
+
+            model = model.from_pretrained(snapshot_download(pretrained_model_dir))
         self.llm = self.register_core_module(model)
         self.learning_rate = learning_rate
         self.use_tril_attention_mask = use_tril_attention_mask

wit/meaning_dataset.py

@@ -115,7 +115,7 @@ class MeaningDataset(Dataset):
 
     def __getitem__(self, idx):
         output = {}
-        data = torch.tensor(self.data[idx])
+        data = torch.tensor(self.data[idx]).long()
         output["input_ids"] = data
         output["labels"] = data.clone()
         output["token_type_ids"] = torch.zeros(data.shape)

wit/special_dataset.py

@@ -10,7 +10,7 @@ from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split, S
 
 
 class SpecialDataset(Dataset):
-    def __init__(self, start=1, end=320, size=32768):  # 1048576 32768
+    def __init__(self, start=1, end=128, size=32768):  # 1048576 32768
         self.size = size
         self.features = []
         a = torch.randint(start, end, [size])
@@ -20,8 +20,8 @@ class SpecialDataset(Dataset):
         z = torch.zeros([size]).long()
         # self.data = torch.stack([a, b, a + b, a + b, a + b * 2]).permute(1, 0)
         # self.data = torch.stack([a, b, a, a + b / 4]).permute(1, 0).long()
-        # self.data = torch.stack([a, a + 1, a + 2]).permute(1, 0).long()
-        self.data = torch.stack([a, b, a]).permute(1, 0).long()
+        self.data = torch.stack([a, a + a, a + a]).permute(1, 0).long()
+        # self.data = torch.stack([a, b, a]).permute(1, 0).long()
         # self.data = torch.stack([a, b, a, a + a / 8, a + a / 4, a + a / 2, a + a]).permute(1, 0).long()
 
         # input  a b c

wit/train.py

@@ -9,50 +9,53 @@ import torch
 from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split, Subset
 
 from transformers import (
-    BatchEncoding,
     DefaultDataCollator,
-    PreTrainedTokenizer,
     set_seed,
 )
-from modelscope import snapshot_download
 from lit_module import LitModule
 from tokenization_qwen import QWenTokenizer
 from logger import TBLogger
 
 from special_dataset import SpecialDataset
 from meaning_dataset import MeaningDataset
+from wit.configuration import ModelConfig
 
-model_name = "qwen/Qwen-1_8B-Chat"
+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"
-tokenizer_name_or_path = None
-train_batch_size = 16
-val_batch_size = 16
+train_batch_size = 256
+val_batch_size = 1
 num_proc = 8
 max_epochs = 1000
 strategy = "auto"
 resume_from_ckpt_path = None
 seed = 42
-vocab_size = 4096
+vocab_size = 256
 
 
 if __name__ == "__main__":
-    if tokenizer_name_or_path is None:
-        tokenizer_name_or_path = model_name
 
     set_seed(seed)
 
-    model_dir = snapshot_download(model_name)
-    lit_module = LitModule(model_dir, learning_rate, use_tril_attention_mask)
+    config = ModelConfig()
+    config.vocab_size = vocab_size
+    config.hidden_size = 128  # 128 1024 2048  32
+    config.num_hidden_layers = 6  # 6 12 24  3
+    config.num_attention_heads = 8  # 8 8 16
+
+    lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
 
     tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
 
-    # raw_dataset = SpecialDataset()
-    raw_dataset = MeaningDataset(start=131072, end=1048576, size=32768)
+    raw_dataset = SpecialDataset()
+    # raw_dataset = MeaningDataset(start=65536, end=262133, size=32768, max_subitem=4, vocab_size=vocab_size)
     train_dataset, val_dataset = random_split(raw_dataset, [0.95, 0.05])
 
-    # daf = next(iter(train_dataset))["input_ids"].numpy().tolist()
+    it = iter(train_dataset)
+    print("data samples:")
+    for i in range(10):
+        print(next(it)["input_ids"].numpy().tolist())
 
     train_dataloader = DataLoader(
         train_dataset,