test_upload / old_model_and_train.py
Chen42's picture
Upload folder using huggingface_hub
41404b8 verified
raw
history blame
8.87 kB
# basic imports
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# transformers imports
from transformers import LiltConfig, BertConfig, EncoderDecoderConfig, EncoderDecoderModel, BertTokenizer, LayoutLMv3Tokenizer, LiltModel
from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
from transformers import default_data_collator
from datasets import load_dataset
# torch imports
import torch
from torch.utils.data import Dataset, DataLoader
# internal imports
# other external imports
import pandas as pd
# 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['block_list']
for box in bboxes:
x0, y0, x1, y1 = box["block_bbox"]
if (x0 > x1) or (y0 > y1):
print(box["block_bbox"])
return False
for cor in box["block_bbox"]:
# if cor < 0 or cor > 1000:
if cor <0:
return False
return True
dataset = load_dataset("json", data_files=data_file)["train"]
print()
print(f"Number of examples: {len(dataset)}")
print()
# print(dataset[0]['block_list'])
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()]
# 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)}")
print(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
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
if model_ckpt_dir:
model.load_state_dict(
torch.load(f"{model_ckpt_dir}/pytorch_model.bin"))
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")
model.load_state_dict(torch.load(f"undertrained/pytorch_model.bin"))
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.
# wc 12420 ./dataset/scene_imgs/jsons/en_json/en_scene.jsonl
# wc 12230 ./dataset/scene_imgs/jsons/zh_json/zh_scene.jsonl
num_instances = 500000
learning_rate = 1e-4
batch_size = 16
num_train_steps = 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/scene_imgs/jsons/en_json/en_scene.jsonl"
data_file = "/home/zychen/hwproject/my_modeling_phase_1/dataset/scene_imgs/jsons/en_json/en_scene.jsonl"
model_ckpt_dir = None
encoder_ckpt_dir = "./Tokenizer_PretrainedWeights/lilt-roberta-en-base"
tgt_tokenizer_dir = "./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")
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,
)
# 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()