120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
from transformers import VisionEncoderDecoderConfig, DonutProcessor, VisionEncoderDecoderModel
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
from pytorch_lightning.loggers import WandbLogger
|
|
from pytorch_lightning.callbacks import ModelCheckpoint
|
|
import pytorch_lightning as pl
|
|
import os
|
|
from huggingface_hub import login
|
|
import argparse
|
|
from sconf import Config
|
|
from utils.checkpoint import CustomCheckpointIO
|
|
from utils.donut_dataset_stream import DonutDataset
|
|
from utils.donut_model_pl import DonutModelPLModule
|
|
from utils.callbacks import PushToHubCallback
|
|
import warnings
|
|
from datasets import load_dataset
|
|
|
|
|
|
|
|
|
|
def main(config, hug_token):
|
|
|
|
config_vision = VisionEncoderDecoderConfig.from_pretrained(
|
|
config.pretrained_model_path)
|
|
config_vision.encoder.image_size = config.image_size
|
|
config_vision.decoder.max_length = config.max_length
|
|
|
|
processor = DonutProcessor.from_pretrained(config.start_model_path)
|
|
model = VisionEncoderDecoderModel.from_pretrained(
|
|
config.pretrained_model_path, config=config_vision)
|
|
|
|
processor.image_processor.size = config.image_size[::-1]
|
|
processor.image_processor.do_align_long_axis = False
|
|
|
|
added_tokens = []
|
|
|
|
dataset = load_dataset(config.dataset_path)
|
|
dataset.train_test_split(test_size=0.1)
|
|
print(dataset)
|
|
|
|
# train_dataset = DonutDataset(
|
|
# dataset,
|
|
# processor=processor,
|
|
# model=model,
|
|
# max_length=config.max_length,
|
|
# split="train",
|
|
# task_start_token="<s_cord-v2>",
|
|
# prompt_end_token="<s_cord-v2>",
|
|
# added_tokens=added_tokens,
|
|
# sort_json_key=False, # cord dataset is preprocessed, so no need for this
|
|
# )
|
|
|
|
# val_dataset = DonutDataset(
|
|
# dataset,
|
|
# processor=processor,
|
|
# model=model,
|
|
# max_length=config.max_length,
|
|
# split="validation",
|
|
# task_start_token="<s_cord-v2>",
|
|
# prompt_end_token="<s_cord-v2>",
|
|
# added_tokens=added_tokens,
|
|
# sort_json_key=False, # cord dataset is preprocessed, so no need for this
|
|
# )
|
|
|
|
# model.config.pad_token_id = processor.tokenizer.pad_token_id
|
|
# model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]
|
|
|
|
# train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=1)
|
|
# val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1)
|
|
|
|
# login(hug_token, True)
|
|
|
|
# model_module = DonutModelPLModule(config.train_config.toDict(), processor, model, max_length=config.max_length, train_dataloader=train_dataloader, val_dataloader=val_dataloader)
|
|
|
|
# wandb_logger = WandbLogger(project="Donut", name=config.wandb_test_name)
|
|
|
|
# checkpoint_callback = ModelCheckpoint(
|
|
# monitor="val_metric",
|
|
# dirpath=config.checkpoint_path,
|
|
# filename="artifacts",
|
|
# save_top_k=1,
|
|
# save_last=False,
|
|
# mode="min",
|
|
# )
|
|
|
|
# custom_ckpt = CustomCheckpointIO()
|
|
|
|
# trainer = pl.Trainer(
|
|
# accelerator="gpu" if torch.cuda.is_available() else 'cpu', # change to gpu
|
|
# devices=1,
|
|
# max_epochs=config.train_config.max_epochs,
|
|
# val_check_interval=config.train_config.val_check_interval,
|
|
# check_val_every_n_epoch=config.train_config.check_val_every_n_epoch,
|
|
# gradient_clip_val=config.train_config.gradient_clip_val,
|
|
# precision=16, # we'll use mixed precision
|
|
# plugins=custom_ckpt,
|
|
# num_sanity_val_steps=0,
|
|
# logger=wandb_logger,
|
|
# callbacks=[PushToHubCallback(output_model_path=config.output_model_path), checkpoint_callback],
|
|
# )
|
|
|
|
# trainer.fit(model_module)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", type=str, required=True)
|
|
args, left_argv = parser.parse_known_args()
|
|
config = Config(args.config)
|
|
config.argv_update(left_argv)
|
|
|
|
hug_token = os.environ.get("HUG_TOKEN", None)
|
|
|
|
if not torch.cuda.is_available():
|
|
warnings.warn("You don't have cuda available, training might be taking long time or impossible")
|
|
|
|
if not hug_token:
|
|
raise Exception("You need to set up HUG_TOKEN in enviroments to push output model to hub")
|
|
main(config, hug_token)
|