Update inference for debug.

wit/inference.py

@@ -1,68 +1,56 @@
-import argparse
-from functools import partial
-from itertools import chain
-from typing import Dict, Tuple
-
-import datasets
 import pytorch_lightning as pl
 import torch
-from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split, Subset
 
-from lit_module import LitModule
-from wit.model.tokenization_qwen import QWenTokenizer
-from logger import TBLogger
-
-from meaning_dataset import MeaningDataset, BatchGroupMeaningDataloader
-from wit.configuration import ModelConfig
-
-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 = 1
-val_batch_size = 2
-num_proc = 8
-max_epochs = 10
-strategy = "auto"
-resume_from_ckpt_path = None
-seed = 42
-vocab_size = 16
+from model.qwen_module import QwenModule
+from model.modeling_wit import QwenRunner
+from model.tokenization_qwen import QWenTokenizer
 
+import configuration
+import dataset.dataset as ds
 
 if __name__ == "__main__":
-    torch.manual_seed(seed)
 
-    config = ModelConfig()
-    config.vocab_size = vocab_size
-    config.hidden_size = 1024  # 128 1024 2048  32
-    config.num_hidden_layers = 1  # 6 12 24  3
+    conf = configuration.TrainConfig()
+    config = conf.model_config
+
+    conf.name = "bigger"  # current train process name
+    conf.pretrain_model_name = None  # "qwen/Qwen-1_8B-Chat"
+    conf.learning_rate = 0.0001
+    conf.use_tril_attention_mask = None
+    conf.precision = "bf16-mixed"  # "precision:bf16-mixed,16-mixed,32-true"
+    conf.train_batch_size = 16
+    conf.val_batch_size = 4
+    conf.num_proc = 8
+    conf.max_epochs = 1000
+    conf.strategy = "auto"
+    conf.resume_from_ckpt_path = None
+    conf.seed = 42
+    conf.dataloader_works = 2
+
+    conf.dataset.meaning.val_mask_level = [0, 1, 2]
+    conf.dataset.meaning.val_mask_idx = [0, 0, -1]
+
+    config.vocab_size = 256
+    config.hidden_size = 128  # 128 1024 2048  32
+    config.num_hidden_layers = 3  # 6 12 24  3
     config.num_attention_heads = 16  # 8 8 16
 
-    lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
-    tokenizer = QWenTokenizer("./model/wit_b64.tiktoken", "./model/wit_char.tiktoken")
+    torch.manual_seed(conf.seed)
 
-    level_ratio = 2
-    start = vocab_size * level_ratio * level_ratio
-    end = start * level_ratio
-    size = end * level_ratio
-    size = 1024
-    raw_dataset = MeaningDataset(start, end, size, vocab_size, level_ratio)
-    train_dataset, val_dataset = raw_dataset.Split(0.95)
+    qwen = QwenModule.load_from_checkpoint(checkpoint_path = "log/bigger/version_1/checkpoints/epoch=26-step=27891.ckpt")
+    qwen.eval()
 
-    train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
-    val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)
+    runner = QwenRunner(qwen.llm)
 
+    train_dataloader, val_dataloader = ds.InitDataset(conf)
     it = iter(val_dataloader)
     batch = next(it)
-    b, l = lit_module.llm(**batch, return_dict=True)
-    print("b ")
-    print(b.detach().cpu().numpy())
 
-    # batch["input_ids"] = batch["input_ids"][0:1, :]
-    batch["input_ids"] = batch["input_ids"][1:2, :]
-    batch["labels"] = batch["labels"][1:2, :]
-    s, l = lit_module.llm(**batch, return_dict=True)
-    print("s ")
-    print(s.detach().cpu().numpy())
+    fdsafd = batch["input_ids"].numpy()
+
+    
+    print(batch["input_ids"].numpy())
+    print(batch["input_ids"][0:1,:-1].numpy())
+    next_token = runner.ChatToken(batch["input_ids"][0:1,:-1].cuda())
+    print(next_token.detach().cpu().numpy())
 
-    print("data samples:")

wit/model/qwen_module.py

@@ -71,6 +71,7 @@ class QwenModule(pl.LightningModule):
     def on_validation_epoch_end(self) -> None:
         self.log("val_loss", self.metric_loss, rank_zero_only=True)
         self.log("accuracy", self.metric_accuracy, rank_zero_only=True)
+        self.log("hp_metric", self.metric_accuracy, rank_zero_only=True)
 
     def configure_optimizers(self):
         optimizer = torch.optim.AdamW(self.trainer.model.parameters(), lr=self.learning_rate)

wit/train.py

@@ -1,8 +1,8 @@
 import pytorch_lightning as pl
 import torch
 
-from model.lit_module import LitModule
-from wit.model.tokenization_qwen import QWenTokenizer
+from model.qwen_module import QwenModule
+from model.tokenization_qwen import QWenTokenizer
 from logger import MLFLogger, TBLogger
 
 import configuration
@@ -27,8 +27,8 @@ if __name__ == "__main__":
     conf.seed = 42
     conf.dataloader_works = 2
 
-    conf.dataset.meaning.mask_level = [0, 1, 2]
-    conf.dataset.meaning.mask_idx = [0, 0, -1]
+    conf.dataset.meaning.val_mask_level = [0, 1, 2]
+    conf.dataset.meaning.val_mask_idx = [0, 0, -1]
 
     config.vocab_size = 256
     config.hidden_size = 128  # 128 1024 2048  32
@@ -36,7 +36,7 @@ if __name__ == "__main__":
     config.num_attention_heads = 16  # 8 8 16
 
     torch.manual_seed(conf.seed)
-    lit_module = LitModule(conf)
+    qwen = QwenModule(conf)
 
     train_dataloader, val_dataloader = ds.InitDataset(conf)
     # for i in range(len(train_dataloader)):
@@ -56,7 +56,7 @@ if __name__ == "__main__":
     )
 
     lit_trainer.fit(
-        lit_module,
+        qwen,
         train_dataloaders=train_dataloader,
         val_dataloaders=val_dataloader,
         ckpt_path=conf.resume_from_ckpt_path,