.gitignore

@@ -2,4 +2,3 @@ __pycache__
 .vscode
 *.txt
 temp
-lightning_logs

dataset/MNBVC.py

@@ -0,0 +1,5 @@
+from datasets import load_dataset
+
+dataset = load_dataset("liwu/MNBVC", "wikipedia", split="train", streaming=True)
+
+print(next(iter(dataset)))  # get the first line

wit/lit_module.py

@@ -1,101 +0,0 @@
-from functools import cache
-from typing import Dict, Optional
-
-import pytorch_lightning as pl
-import torch
-import torchmetrics
-
-# from utils import init_model
-# from custom_models.gpt2.modeling_gpt2 import GPT2LMHeadModel
-
-from modeling_wit import QWenLMHeadModel
-from configuration_qwen import QWenConfig
-
-from transformers import AutoConfig
-
-
-class LitModule(pl.LightningModule):
-    def __init__(
-        self,
-        model_dir: str,
-        learning_rate: float = 0.0001,
-        use_tril_attention_mask: str = False,
-    ):
-        super().__init__()
-        self.save_hyperparameters()
-        config = QWenConfig()
-        model = QWenLMHeadModel(config)
-        model = model.from_pretrained(model_dir)
-        self.llm = self.register_core_module(model)
-        self.learning_rate = learning_rate
-        self.use_tril_attention_mask = use_tril_attention_mask
-        self.metric_loss = torchmetrics.MeanMetric()
-        self.metric_accuracy = torchmetrics.Accuracy(
-            task="multiclass",
-            num_classes=self.llm.config.vocab_size,
-        )
-
-    @cache
-    def get_batch_tril_matrix(self, block_size: int, batch_size: Optional[int] = None) -> torch.Tensor:
-        matrix = torch.ones(block_size, block_size).tril()
-        if batch_size is not None:
-            matrix = matrix.repeat(batch_size, 1, 1)
-        return matrix
-
-    def register_core_module(self, module: torch.nn.Module) -> torch.nn.Module:
-        object.__setattr__(self, "__core_module__", module)
-        return module
-
-    def training_step(self, batch: Dict[str, torch.Tensor], batch_idx):
-        batch_size, block_size = batch["input_ids"].shape
-        if self.use_tril_attention_mask:
-            batch["attention_mask"] = self.get_batch_tril_matrix(block_size, batch_size=batch_size).to(self.device)
-        outputs, loss = self.llm(**batch)
-        self.log("train_loss", loss, rank_zero_only=True)
-        return loss
-
-    def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx):
-        outputs, loss = self.llm(**batch, return_dict=True)
-        logits = outputs[..., :-1, :]
-        labels = batch["labels"][..., 1:]
-
-        self.metric_loss.update(loss)
-
-        label_mask = labels != -100
-        self.metric_accuracy.update(logits[label_mask], labels[label_mask])
-
-    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)
-
-    def configure_optimizers(self):
-        strategy = self.trainer.strategy
-        if isinstance(strategy, pl.strategies.DeepSpeedStrategy):
-            assert "optimizer" not in strategy.config
-            zero_config = strategy.config.get("zero_optimization")
-            if zero_config is not None:
-                if "offload_optimizer" in zero_config:
-                    import deepspeed
-
-                    optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam(
-                        self.trainer.model.parameters(), lr=self.learning_rate
-                    )
-                    return optimizer
-        optimizer = torch.optim.AdamW(self.trainer.model.parameters(), lr=self.learning_rate)
-        return optimizer
-
-    def configure_callbacks(self):
-        checkpoint_callback = pl.callbacks.ModelCheckpoint(
-            monitor="accuracy",
-            mode="max",
-            filename="{epoch:02d}-{accuracy:.4f}",
-        )
-        early_stop_callback = pl.callbacks.EarlyStopping(
-            monitor="accuracy",
-            min_delta=0.001,
-            patience=3,
-            mode="max",
-            stopping_threshold=1,
-        )
-        return [checkpoint_callback]
-        # return [checkpoint_callback, early_stop_callback]

wit/lit_train.py

@@ -1,181 +0,0 @@
-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
-from transformers import (
-    BatchEncoding,
-    DefaultDataCollator,
-    PreTrainedTokenizer,
-    set_seed,
-)
-from modelscope import snapshot_download
-from lit_module import LitModule
-from tokenization_qwen import QWenTokenizer
-
-model_name = "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
-dataset_name = ["/home/colin/develop/dataset/liwu/MNBVC/wiki"]
-dataset_name = ["/home/colin/develop/dataset/liwu/MNBVC/wiki/20230198/58.jsonl.gz"]
-train_batch_size = 256
-val_batch_size = 16
-limit_val_batches = 128
-num_proc = 8
-max_epochs = 1000
-strategy = "fsdp"
-resume_from_ckpt_path = None
-seed = 42
-
-
-class SpecialDataset(Dataset):
-    def __init__(self, size=65536):
-        self.size = size
-        self.features = []
-        a = torch.randint(0, 1024, [size])
-        self.data = torch.stack([a, a * 2, a * 3, a * 4]).permute(1, 0)
-
-    def __len__(self):
-        return self.size
-
-    def __getitem__(self, idx):
-        output = {}
-        data = self.data[idx]
-        output["input_ids"] = data
-        output["labels"] = data
-        output["token_type_ids"] = torch.zeros(data.shape)
-        return output
-
-
-def split_raw_dataset(
-    raw_dataset: datasets.DatasetDict,
-) -> Tuple[datasets.Dataset, datasets.Dataset]:
-    if "validation" in raw_dataset:
-        train_dataset, val_dataset = raw_dataset["train"], raw_dataset["validation"]
-    else:
-        raw_dataset = raw_dataset["train"].train_test_split(test_size=0.05, seed=seed)
-        train_dataset, val_dataset = raw_dataset["train"], raw_dataset["test"]
-    return train_dataset, val_dataset
-
-
-def process_dataset(dataset: datasets.Dataset, tokenizer: PreTrainedTokenizer) -> datasets.Dataset:
-    def group_texts(examples: Dict[str, list], block_size: int = 512) -> BatchEncoding:
-        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
-        total_length = len(concatenated_examples[list(examples.keys())[0]])
-        total_length = (total_length // block_size) * block_size
-        result = {
-            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
-            for k, t in concatenated_examples.items()
-        }
-        result["labels"] = result["input_ids"].copy()
-        result = BatchEncoding(result)
-        return result
-
-    def format_inputs(examples):
-        p = examples["段落"]
-        mergeLine = ""
-        for line in p:
-            mergeLine += line["内容"] + "\n"
-        return {"text": mergeLine}
-
-    def tokenize_inputs(
-        examples: Dict[str, list],
-        tokenizer: PreTrainedTokenizer,
-        column_name: str = "text",
-    ) -> BatchEncoding:
-        logits = tokenizer(examples[column_name], return_attention_mask=False)
-        return logits
-
-    dataset_column_names = list(dataset.features)
-    dataset = dataset.map(
-        partial(format_inputs),
-        batched=False,
-        num_proc=num_proc,
-        remove_columns=dataset_column_names,
-    )
-    dataset_column_names = list(dataset.features)
-    dataset = dataset.map(
-        partial(tokenize_inputs, tokenizer=tokenizer),
-        batched=True,
-        num_proc=num_proc,
-        remove_columns=dataset_column_names,
-    )
-    dataset = dataset.map(
-        partial(group_texts, block_size=tokenizer.model_max_length),
-        batched=True,
-        num_proc=num_proc,
-    )
-
-    return dataset
-
-
-if __name__ == "__main__":
-    if tokenizer_name_or_path is None:
-        tokenizer_name_or_path = model_name
-
-    set_seed(seed)
-
-    # lightning module
-    model_dir = snapshot_download(model_name)
-    lit_module = LitModule(model_dir, learning_rate, use_tril_attention_mask)
-
-    tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
-    train_dataset_list = []
-    val_dataset_list = []
-    for dn in dataset_name:
-        datanames = dn.split(".")
-        if datanames[-1] == "gz" and datanames[-2] == "jsonl":
-            raw_dataset = datasets.load_dataset("json", data_files=dn)
-        elif datanames[-1] == "json":
-            raw_dataset = datasets.load_dataset("json", data_files=dn)
-        else:
-            raw_dataset = datasets.load_dataset(dn)
-        train_dataset, val_dataset = split_raw_dataset(raw_dataset)
-        train_dataset = process_dataset(train_dataset, tokenizer)
-        val_dataset = process_dataset(val_dataset, tokenizer)
-        train_dataset_list.append(train_dataset)
-        val_dataset_list.append(val_dataset)
-    train_dataset = ConcatDataset(train_dataset_list)
-    val_dataset = ConcatDataset(val_dataset_list)
-
-    train_dataset = SpecialDataset()
-    val_dataset = SpecialDataset()
-
-    train_dataloader = DataLoader(
-        train_dataset,
-        batch_size=train_batch_size,
-        num_workers=num_proc,
-        collate_fn=DefaultDataCollator(),
-        persistent_workers=True,
-        shuffle=True,
-    )
-    val_dataloader = DataLoader(
-        val_dataset,
-        batch_size=val_batch_size,
-        num_workers=num_proc,
-        collate_fn=DefaultDataCollator(),
-        persistent_workers=True,
-        shuffle=True,
-    )
-
-    torch.set_float32_matmul_precision("medium")
-    precision = precision
-    lit_trainer = pl.Trainer(
-        accelerator="gpu",
-        precision=precision,
-        strategy=strategy,
-        max_epochs=max_epochs,
-        limit_val_batches=limit_val_batches,
-    )
-    lit_trainer.fit(
-        lit_module,
-        train_dataloaders=train_dataloader,
-        val_dataloaders=val_dataloader,
-        ckpt_path=resume_from_ckpt_path,
-    )

wit/modeling_wit.py

@@ -137,16 +137,6 @@ class QWenLMHeadModel(nn.Module):
         self.transformer = QWenModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        **kwargs,
-    ):
-        runner = QwenRunner(self)
-        return runner.forwardQWen(input_ids, labels)
-
     def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]]):
         pretrained_model_name_or_path = str(pretrained_model_name_or_path)
         resolved_archive_file = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json")
@@ -353,7 +343,15 @@ class QwenRunner:
             loss_fct = CrossEntropyLoss()
             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
 
-        return lm_logits, loss
+        # shift_labels = torch.ones([1,19]).to(lm_logits.device).to(torch.int64)
+        # shift_logits = lm_logits[..., :-1, :].contiguous()
+        # loss_fct = CrossEntropyLoss()
+        # loss = loss_fct(
+        #     shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+        # )
+        # loss.backward()
+
+        return lm_logits
 
     def prepareInput(self, tokenizer, query, query_assistant, history, system):
         return make_context(tokenizer, query, query_assistant, history=history, system=system)

wit/tokenization_qwen.py

@@ -63,7 +63,7 @@ class QWenTokenizer(PreTrainedTokenizer):
 
         self.mergeable_ranks = _load_tiktoken_b64(vocab_file_b64)
         self.mergeable_ranks.update(_load_tiktoken_char(vocab_file_char, len(self.mergeable_ranks)))
-        self.model_max_length = 1024
+
         special = (
             "user",
             "assistant",