model_and_train.py

# basic imports
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "4"

# other external imports
import pandas as pd
# torch imports
import torch
from datasets import load_dataset
from torch.utils.data import DataLoader, Dataset
# transformers imports
from transformers import (BertConfig, BertTokenizer, EncoderDecoderConfig,
                          EncoderDecoderModel, LayoutLMv3Tokenizer, LiltConfig,
                          LiltModel, Seq2SeqTrainer, Seq2SeqTrainingArguments,
                          default_data_collator)

# internal imports



# prepare tokenizer.
def prepare_tokenizer(src_tokenizer_dir, tgt_tokenizer_dir):
    src_tokenizer = LayoutLMv3Tokenizer.from_pretrained(src_tokenizer_dir)
    tgt_tokenizer = BertTokenizer.from_pretrained(tgt_tokenizer_dir)

    return src_tokenizer, tgt_tokenizer


# read data points.
def prepare_dataset_df(data_file):

    def filter_fn(exam):
        bboxes = exam["layout_src"]
        for box in bboxes:
            x0, y0, x1, y1 = box
            if (x0 > x1) or (y0 > y1):
                print("(x0 > x1) or (y0 > y1)")
                return False
            for cor in box:
                if cor < 0 or cor > 1000:
                    # print("cor < 0 or cor > 1000")
                    # print(exam['img_path'],box)
                    return False
        return True

    dataset = load_dataset("json", data_files=data_file)["train"]
    print()
    print(f"Number of examples: {len(dataset)}")
    print()

    dataset = dataset.filter(filter_fn, num_proc=48)

    dataset_df = dataset.to_pandas()
    # dataset_df = pd.read_json(data_file, lines=True, orient="records")

    # filter the nan data points.
    dataset_df = dataset_df[~dataset_df["tgt_sen_trans"].isna()]
    dataset_df = dataset_df[~dataset_df["text_src"].isna()]
    dataset_df = dataset_df[~dataset_df["layout_src"].isna()]
    # remove entries where "text_src" length is less than 3
    dataset_df = dataset_df[dataset_df["text_src"].str.len() >= 3]
    # reconstruct the idx to avoid index_error.
    dataset_df = dataset_df.reset_index(drop=True)

    print(f"Number of examples after filtered: {len(dataset_df)}")
    return dataset_df


class MyDataset(Dataset):

    def __init__(
        self,
        df,
        src_tokenizer,
        tgt_tokenizer,
        max_src_length,
        max_target_length,
    ):
        self.df = df
        self.src_tokenizer = src_tokenizer
        self.tgt_tokenizer = tgt_tokenizer
        self.max_src_length = max_src_length
        self.max_target_length = max_target_length

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # get text_src + layout_src + tgt_trans.
        text_src = self.df['text_src'][idx]
        layout_src = self.df['layout_src'][idx]
        tgt_trans = self.df['tgt_sen_trans'][idx]

        # read in annotations at word-level (words, word boxes)
        words_ = text_src.split(" ")
        word_boxes_ = layout_src
        # print('words', words_, len(words_), len(word_boxes_))
        assert len(words_) == len(word_boxes_)
        words = []
        word_boxes = []
        for word, word_box in zip(words_, word_boxes_):
            if (word_box[0] >= word_box[2]) or (word_box[1] >= word_box[3]):
                continue

            words.append(word)
            word_boxes.append(word_box)

        assert len(words) == len(word_boxes)

        encoding = self.src_tokenizer(
            words,
            boxes=word_boxes,
            padding="max_length",
            truncation=True,
            max_length=self.max_src_length,
        )

        # construct labels.
        labels = self.tgt_tokenizer(
            tgt_trans,
            padding="max_length",
            truncation=True,
            max_length=self.max_target_length)["input_ids"]
        # important: make sure that PAD tokens are ignored by the loss function
        labels = [
            label if label != self.tgt_tokenizer.pad_token_id else -100
            for label in labels
        ]

        encoding["labels"] = labels

        assert len(encoding['input_ids']) == self.max_src_length
        assert len(encoding['attention_mask']) == self.max_src_length
        assert len(encoding['bbox']) == self.max_src_length
        assert len(encoding['labels']) == self.max_target_length

        # finally, convert everything to PyTorch tensors
        for k, v in encoding.items():
            encoding[k] = torch.as_tensor(encoding[k])

        return encoding


def prepare_model(src_tokenizer,
                  tgt_tokenizer,
                  max_src_len,
                  max_tgt_len,
                  num_encoder_hidden_layers,
                  num_decoder_hidden_layers,
                  encoder_ckpt_dir,
                  model_ckpt_dir=None):
    config_encoder = LiltConfig.from_pretrained(
        encoder_ckpt_dir,
        max_position_embeddings=max_src_len + 2,
        num_hidden_layers=num_encoder_hidden_layers)
    config_decoder = BertConfig(vocab_size=tgt_tokenizer.vocab_size,
                                max_position_embeddings=max_tgt_len,
                                num_hidden_layers=num_decoder_hidden_layers)

    model_config = EncoderDecoderConfig.from_encoder_decoder_configs(
        encoder_config=config_encoder,
        decoder_config=config_decoder,
    )
    model = EncoderDecoderModel(config=model_config, )

    model.config.decoder_start_token_id = tgt_tokenizer.cls_token_id
    model.config.pad_token_id = tgt_tokenizer.pad_token_id
    model.config.vocab_size = tgt_tokenizer.vocab_size
    model.config.eos_token_id = tgt_tokenizer.pad_token_id

    from safetensors.torch import load_file
    if model_ckpt_dir:
        bin_path = f"{model_ckpt_dir}/pytorch_model.bin"
        safetensors_path = f"{model_ckpt_dir}/model.safetensors"
        if os.path.exists(bin_path):
            state_dict = torch.load(bin_path)
        elif os.path.exists(safetensors_path):
            state_dict = load_file(safetensors_path)
        else:
            raise FileNotFoundError(
                "Neither pytorch_model.bin nor model.safetensors found in the specified directory."
            )
        model.load_state_dict(state_dict, strict=False)
        model.save_pretrained(
            f"continued_{model_ckpt_dir}")  #save at continued training
    else:
        # Loading the pre-trained params and then save the model, including its configuration.
        tmp_encoder = LiltModel.from_pretrained(
            pretrained_model_name_or_path=encoder_ckpt_dir,
            config=config_encoder,
        )
        # tmp_encoder = LiltModel(config=config_encoder)
        model.encoder = tmp_encoder
        # model.save_pretrained("undertrained_default_safe_true")
        model.save_pretrained("undertrained_safe_serialization_False", safe_serialization=False)
        # model.load_state_dict(torch.load(f"undertrained/pytorch_model.bin"))

        bin_path = "undertrained_safe_serialization_False/pytorch_model.bin"
        safetensors_path = "undertrained_default_safe_true/model.safetensors"
        if os.path.exists(bin_path):
            state_dict = torch.load(bin_path)
        elif os.path.exists(safetensors_path):
            state_dict = load_file(safetensors_path)
        else:
            raise FileNotFoundError(
                "Neither pytorch_model.bin nor model.safetensors found in the specified directory."
            )
        model.load_state_dict(state_dict, strict=False)

    print(model.config)
    print(model)

    return model


if __name__ == "__main__":

    # hyper-parameters.
    ## for model.
    MAX_TGT_LEN = 512
    MAX_SRC_LEN = 512
    num_encoder_hidden_layers = 12
    num_decoder_hidden_layers = 12

    ## for training.
    num_instances = 500000  #total 620082 ./dataset/merged.jsonl Number of examples after filtered: 547084
    learning_rate = 1e-4
    batch_size = 28
    num_train_steps = 400000 #400000
    output_dir = f"./train.lr_{learning_rate}.bsz_{batch_size}.step_{num_train_steps}.layer_{num_encoder_hidden_layers}-{num_decoder_hidden_layers}"
    save_total_limit = 100
    save_steps = num_train_steps // save_total_limit

    dataset_dir = "/home/zychen/hwproject/my_modeling_phase_1/dataset"
    data_file = f"{dataset_dir}/merged.jsonl"

    # model_ckpt_dir = '/home/zychen/hwproject/my_modeling_phase_1/train.lr_0.0001.bsz_8.step_400000.layer_12-12/checkpoint-32000'
    model_ckpt_dir = '/home/zychen/hwproject/my_modeling_phase_1/train.lr_0.0001.bsz_16.step_500000.layer_12-12_36k+20k/checkpoint-20000'
    encoder_ckpt_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/lilt-roberta-en-base"

    tgt_tokenizer_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/bert-base-chinese-tokenizer"

    src_tokenizer, tgt_tokenizer = prepare_tokenizer(
        src_tokenizer_dir=encoder_ckpt_dir,
        tgt_tokenizer_dir=tgt_tokenizer_dir,
    )
    dataset_df = prepare_dataset_df(data_file=data_file)[:num_instances]
    print(f"\nnum_instances: {len(dataset_df)}\n")
    print(dataset_df)
    my_dataset = MyDataset(
        df=dataset_df,
        src_tokenizer=src_tokenizer,
        tgt_tokenizer=tgt_tokenizer,
        max_src_length=MAX_SRC_LEN,
        max_target_length=MAX_TGT_LEN,
    )
    model = prepare_model(src_tokenizer=src_tokenizer,
                          tgt_tokenizer=tgt_tokenizer,
                          max_src_len=MAX_SRC_LEN,
                          max_tgt_len=MAX_TGT_LEN,
                          num_encoder_hidden_layers=num_encoder_hidden_layers,
                          num_decoder_hidden_layers=num_decoder_hidden_layers,
                          encoder_ckpt_dir=encoder_ckpt_dir,
                          model_ckpt_dir=model_ckpt_dir)

    training_args = Seq2SeqTrainingArguments(
        predict_with_generate=False,
        evaluation_strategy="no",
        per_device_train_batch_size=batch_size,
        fp16=True,
        output_dir=output_dir,
        logging_steps=1,
        # save_strategy="epoch",
        learning_rate=learning_rate,
        max_steps=num_train_steps,
        warmup_ratio=0.05,
        save_total_limit=save_total_limit,
        save_steps=save_steps,
        save_safetensors=False,
    )
    # print(training_args)
    # instantiate trainer
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        compute_metrics=None,
        train_dataset=my_dataset,
        eval_dataset=None,
        data_collator=default_data_collator,
    )

    trainer.train()