streaming
This commit is contained in:
parent
ab051e21b1
commit
679048f88a
115
train_stream.py
Normal file
115
train_stream.py
Normal file
@ -0,0 +1,115 @@
|
||||
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 = []
|
||||
|
||||
train_dataset = DonutDataset(
|
||||
config.dataset_path,
|
||||
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(
|
||||
config.dataset_path,
|
||||
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=4)
|
||||
val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=4)
|
||||
|
||||
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)
|
155
utils/donut_dataset_stream.py
Normal file
155
utils/donut_dataset_stream.py
Normal file
@ -0,0 +1,155 @@
|
||||
from datasets import load_dataset
|
||||
from torch.utils.data import Dataset
|
||||
import json
|
||||
from typing import Any, List, Tuple
|
||||
import random
|
||||
import torch
|
||||
from transformers import DonutProcessor, VisionEncoderDecoderModel
|
||||
|
||||
|
||||
class DonutDataset(Dataset):
|
||||
"""
|
||||
DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets)
|
||||
Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),
|
||||
and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string).
|
||||
Args:
|
||||
dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl
|
||||
max_length: the max number of tokens for the target sequences
|
||||
split: whether to load "train", "validation" or "test" split
|
||||
ignore_id: ignore_index for torch.nn.CrossEntropyLoss
|
||||
task_start_token: the special token to be fed to the decoder to conduct the target task
|
||||
prompt_end_token: the special token at the end of the sequences
|
||||
sort_json_key: whether or not to sort the JSON keys
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_name_or_path: str,
|
||||
max_length: int,
|
||||
processor: DonutProcessor,
|
||||
model: VisionEncoderDecoderModel,
|
||||
split: str = "train",
|
||||
ignore_id: int = -100,
|
||||
task_start_token: str = "<s>",
|
||||
prompt_end_token: str = None,
|
||||
sort_json_key: bool = True,
|
||||
added_tokens: list = []
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.max_length = max_length
|
||||
self.split = split
|
||||
self.processor = processor
|
||||
self.model = model
|
||||
self.ignore_id = ignore_id
|
||||
self.task_start_token = task_start_token
|
||||
self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token
|
||||
self.sort_json_key = sort_json_key
|
||||
self.added_tokens = added_tokens
|
||||
|
||||
self.dataset = load_dataset(dataset_name_or_path, split=self.split, stream=True).with_format("torch")
|
||||
self.dataset_length = len(self.dataset)
|
||||
|
||||
self.gt_token_sequences = []
|
||||
for sample in self.dataset:
|
||||
ground_truth = json.loads(sample["ground_truth"])
|
||||
if "gt_parses" in ground_truth: # when multiple ground truths are available, e.g., docvqa
|
||||
assert isinstance(ground_truth["gt_parses"], list)
|
||||
gt_jsons = ground_truth["gt_parses"]
|
||||
else:
|
||||
assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
|
||||
gt_jsons = [ground_truth["gt_parse"]]
|
||||
|
||||
self.gt_token_sequences.append(
|
||||
[
|
||||
self.json2token(
|
||||
gt_json,
|
||||
update_special_tokens_for_json_key=self.split == "train",
|
||||
sort_json_key=self.sort_json_key,
|
||||
)
|
||||
+ self.processor.tokenizer.eos_token
|
||||
for gt_json in gt_jsons # load json from list of json
|
||||
]
|
||||
)
|
||||
|
||||
self.add_tokens([self.task_start_token, self.prompt_end_token])
|
||||
self.prompt_end_token_id = self.processor.tokenizer.convert_tokens_to_ids(self.prompt_end_token)
|
||||
|
||||
def json2token(self, obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True):
|
||||
"""
|
||||
Convert an ordered JSON object into a token sequence
|
||||
"""
|
||||
if type(obj) == dict:
|
||||
if len(obj) == 1 and "text_sequence" in obj:
|
||||
return obj["text_sequence"]
|
||||
else:
|
||||
output = ""
|
||||
if sort_json_key:
|
||||
keys = sorted(obj.keys(), reverse=True)
|
||||
else:
|
||||
keys = obj.keys()
|
||||
for k in keys:
|
||||
if update_special_tokens_for_json_key:
|
||||
self.add_tokens([fr"<s_{k}>", fr"</s_{k}>"])
|
||||
output += (
|
||||
fr"<s_{k}>"
|
||||
+ self.json2token(obj[k], update_special_tokens_for_json_key, sort_json_key)
|
||||
+ fr"</s_{k}>"
|
||||
)
|
||||
return output
|
||||
elif type(obj) == list:
|
||||
return r"<sep/>".join(
|
||||
[self.json2token(item, update_special_tokens_for_json_key, sort_json_key) for item in obj]
|
||||
)
|
||||
else:
|
||||
obj = str(obj)
|
||||
if f"<{obj}/>" in self.added_tokens:
|
||||
obj = f"<{obj}/>" # for categorical special tokens
|
||||
return obj
|
||||
|
||||
def add_tokens(self, list_of_tokens: List[str]):
|
||||
"""
|
||||
Add special tokens to tokenizer and resize the token embeddings of the decoder
|
||||
"""
|
||||
newly_added_num = self.processor.tokenizer.add_tokens(list_of_tokens)
|
||||
if newly_added_num > 0:
|
||||
self.model.decoder.resize_token_embeddings(len(self.processor.tokenizer))
|
||||
self.added_tokens.extend(list_of_tokens)
|
||||
|
||||
# def __len__(self) -> int:
|
||||
# return self.dataset_length
|
||||
|
||||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Load image from image_path of given dataset_path and convert into input_tensor and labels
|
||||
Convert gt data into input_ids (tokenized string)
|
||||
Returns:
|
||||
input_tensor : preprocessed image
|
||||
input_ids : tokenized gt_data
|
||||
labels : masked labels (model doesn't need to predict prompt and pad token)
|
||||
"""
|
||||
sample = self.dataset[idx]
|
||||
|
||||
# change if not 3 channels
|
||||
if sample['image'].mode != "RGB":
|
||||
sample['image'] = sample['image'].convert("RGB")
|
||||
|
||||
# inputs
|
||||
pixel_values = self.processor(sample["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values
|
||||
pixel_values = pixel_values.squeeze()
|
||||
|
||||
# targets
|
||||
target_sequence = random.choice(self.gt_token_sequences[idx]) # can be more than one, e.g., DocVQA Task 1
|
||||
input_ids = self.processor.tokenizer(
|
||||
target_sequence,
|
||||
add_special_tokens=False,
|
||||
max_length=self.max_length,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)["input_ids"].squeeze(0)
|
||||
|
||||
labels = input_ids.clone()
|
||||
labels[labels == self.processor.tokenizer.pad_token_id] = self.ignore_id # model doesn't need to predict pad token
|
||||
# labels[: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1] = self.ignore_id # model doesn't need to predict prompt (for VQA)
|
||||
return pixel_values, labels, target_sequence
|
Loading…
Reference in New Issue
Block a user