70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
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 tokenization_qwen import QWenTokenizer
|
|
from logger import TBLogger
|
|
|
|
from special_dataset import SpecialDataset
|
|
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
|
|
|
|
|
|
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
|
|
config.num_attention_heads = 16 # 8 8 16
|
|
|
|
lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
|
|
tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
|
|
|
|
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)
|
|
|
|
train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
|
|
val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)
|
|
|
|
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())
|
|
|
|
print("data samples:")
|