|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
Fine-tuning the library vision-encoder-decoder models for image captioning. |
|
""" |
|
|
|
import json |
|
import logging |
|
import os |
|
import sys |
|
import time |
|
from dataclasses import asdict, dataclass, field |
|
from enum import Enum |
|
from functools import partial |
|
from pathlib import Path |
|
from typing import Callable, Optional |
|
|
|
import datasets |
|
import nltk |
|
import numpy as np |
|
from datasets import Dataset, load_dataset, load_metric |
|
from tqdm import tqdm |
|
from PIL import Image |
|
|
|
import jax |
|
import jax.numpy as jnp |
|
import optax |
|
import transformers |
|
from filelock import FileLock |
|
from flax import jax_utils, traverse_util |
|
from flax.jax_utils import unreplicate |
|
from flax.training import train_state |
|
from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key |
|
from huggingface_hub import Repository |
|
from transformers import ( |
|
CONFIG_MAPPING, |
|
FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING, |
|
FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, |
|
AutoConfig, |
|
AutoFeatureExtractor, |
|
AutoTokenizer, |
|
HfArgumentParser, |
|
is_tensorboard_available, |
|
FlaxAutoModelForVision2Seq, |
|
) |
|
from transformers.file_utils import get_full_repo_name, is_offline_mode |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
try: |
|
nltk.data.find("tokenizers/punkt") |
|
except (LookupError, OSError): |
|
if is_offline_mode(): |
|
raise LookupError( |
|
"Offline mode: run this script without TRANSFORMERS_OFFLINE first to download nltk data files" |
|
) |
|
with FileLock(".lock") as lock: |
|
nltk.download("punkt", quiet=True) |
|
|
|
|
|
MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING.keys()) |
|
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES) |
|
DECODER_MODEL_TYPES = tuple(conf.model_type for conf in list(FLAX_MODEL_FOR_CAUSAL_LM_MAPPING.keys())) |
|
|
|
|
|
|
|
def shift_tokens_right(input_ids: np.ndarray, pad_token_id: int, decoder_start_token_id: int) -> np.ndarray: |
|
""" |
|
Shift input ids one token to the right. |
|
""" |
|
shifted_input_ids = np.zeros_like(input_ids) |
|
shifted_input_ids[:, 1:] = input_ids[:, :-1] |
|
shifted_input_ids[:, 0] = decoder_start_token_id |
|
|
|
shifted_input_ids = np.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids) |
|
return shifted_input_ids |
|
|
|
|
|
@dataclass |
|
class TrainingArguments: |
|
output_dir: str = field( |
|
metadata={"help": "The output directory where the model predictions and checkpoints will be written."}, |
|
) |
|
overwrite_output_dir: bool = field( |
|
default=False, |
|
metadata={ |
|
"help": ( |
|
"Overwrite the content of the output directory. " |
|
"Use this to continue training if output_dir points to a checkpoint directory." |
|
) |
|
}, |
|
) |
|
do_train: bool = field(default=False, metadata={"help": "Whether to run training."}) |
|
do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the dev set."}) |
|
do_predict: bool = field(default=False, metadata={"help": "Whether to run predictions on the test set."}) |
|
per_device_train_batch_size: int = field( |
|
default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for training."} |
|
) |
|
per_device_eval_batch_size: int = field( |
|
default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for evaluation."} |
|
) |
|
_block_size_doc = \ |
|
""" |
|
The default value `0` will preprocess (tokenization + feature extraction) the whole dataset before training and |
|
cache the results. This uses more disk space, but avoids (repeated) processing time during training. This is a |
|
good option if your disk space is large enough to store the whole processed dataset. |
|
If a positive value is given, the captions in the dataset will be tokenized before training and the results are |
|
cached. During training, it iterates the dataset in chunks of size `block_size`. On each block, images are |
|
transformed by the feature extractor with the results being kept in memory (no cache), and batches of size |
|
`batch_size` are yielded before processing the next block. This could avoid the heavy disk usage when the |
|
dataset is large. |
|
""" |
|
block_size: int = field( |
|
default=0, |
|
metadata={"help": _block_size_doc} |
|
) |
|
learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."}) |
|
weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."}) |
|
adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"}) |
|
adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for AdamW optimizer"}) |
|
adam_epsilon: float = field(default=1e-8, metadata={"help": "Epsilon for AdamW optimizer."}) |
|
label_smoothing_factor: float = field( |
|
default=0.0, metadata={"help": "The label smoothing epsilon to apply (zero means no label smoothing)."} |
|
) |
|
num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."}) |
|
warmup_steps: int = field(default=0, metadata={"help": "Linear warmup over warmup_steps."}) |
|
logging_steps: int = field(default=500, metadata={"help": "Log every X updates steps."}) |
|
eval_steps: int = field(default=None, metadata={"help": "Run an evaluation every X steps."}) |
|
seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."}) |
|
push_to_hub: bool = field( |
|
default=False, metadata={"help": "Whether or not to upload the trained model to the model hub after training."} |
|
) |
|
hub_model_id: str = field( |
|
default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."} |
|
) |
|
hub_token: str = field(default=None, metadata={"help": "The token to use to push to the Model Hub."}) |
|
|
|
def __post_init__(self): |
|
if self.output_dir is not None: |
|
self.output_dir = os.path.expanduser(self.output_dir) |
|
|
|
def to_dict(self): |
|
""" |
|
Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates |
|
the token values by removing their value. |
|
""" |
|
d = asdict(self) |
|
for k, v in d.items(): |
|
if isinstance(v, Enum): |
|
d[k] = v.value |
|
if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum): |
|
d[k] = [x.value for x in v] |
|
if k.endswith("_token"): |
|
d[k] = f"<{k.upper()}>" |
|
return d |
|
|
|
|
|
@dataclass |
|
class ModelArguments: |
|
""" |
|
Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch. |
|
""" |
|
|
|
model_name_or_path: Optional[str] = field( |
|
default=None, |
|
metadata={ |
|
"help": "The model checkpoint for weights initialization." |
|
"Don't set if you want to train a model from scratch." |
|
}, |
|
) |
|
encoder_model_name_or_path: Optional[str] = field( |
|
default=None, |
|
metadata={ |
|
"help": "The encoder model checkpoint for weights initialization." |
|
"Don't set if you want to train an encoder model from scratch." |
|
}, |
|
) |
|
decoder_model_name_or_path: Optional[str] = field( |
|
default=None, |
|
metadata={ |
|
"help": "The decoder model checkpoint for weights initialization." |
|
"Don't set if you want to train a decoder model from scratch." |
|
}, |
|
) |
|
model_type: Optional[str] = field( |
|
default='vision-encoder-decoder', |
|
metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)} |
|
) |
|
encoder_model_type: Optional[str] = field( |
|
default=None, |
|
metadata={"help": "If training from scratch, pass a vision encoder model type from the library. For example, 'vit'"} |
|
) |
|
decoder_model_type: Optional[str] = field( |
|
default=None, |
|
metadata={"help": "If training from scratch, pass a decoder model type from the list: " + ", ".join(DECODER_MODEL_TYPES)} |
|
) |
|
config_name: Optional[str] = field( |
|
default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"} |
|
) |
|
encoder_config_name: Optional[str] = field( |
|
default=None, metadata={"help": "Pretrained encoder config name or path if not the same as encoder_model_name"} |
|
) |
|
decoder_config_name: Optional[str] = field( |
|
default=None, metadata={"help": "Pretrained decoder config name or path if not the same as decoder_model_name"} |
|
) |
|
feature_extractor_name: Optional[str] = field( |
|
default=None, metadata={"help": "Pretrained encoder feature extractor_name or path if not the same as encoder_model_name"} |
|
) |
|
tokenizer_name: Optional[str] = field( |
|
default=None, metadata={"help": "Pretrained decoder tokenizer name or path if not the same as decoder_model_name"} |
|
) |
|
cache_dir: Optional[str] = field( |
|
default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"} |
|
) |
|
use_fast_tokenizer: bool = field( |
|
default=True, |
|
metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."}, |
|
) |
|
dtype: Optional[str] = field( |
|
default="float32", |
|
metadata={ |
|
"help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`." |
|
}, |
|
) |
|
|
|
|
|
@dataclass |
|
class DataTrainingArguments: |
|
""" |
|
Arguments pertaining to what data we are going to input our model for training and eval. |
|
""" |
|
|
|
dataset_name: Optional[str] = field( |
|
default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} |
|
) |
|
dataset_config_name: Optional[str] = field( |
|
default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} |
|
) |
|
data_dir: Optional[str] = field( |
|
default=None, metadata={"help": "The data directory of the dataset to use (via the datasets library)."} |
|
) |
|
image_column: Optional[str] = field( |
|
default=None, |
|
metadata={"help": "The name of the column in the datasets containing the full image file paths."}, |
|
) |
|
caption_column: Optional[str] = field( |
|
default=None, |
|
metadata={"help": "The name of the column in the datasets containing the image captions."}, |
|
) |
|
train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."}) |
|
validation_file: Optional[str] = field( |
|
default=None, |
|
metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."}, |
|
) |
|
test_file: Optional[str] = field( |
|
default=None, |
|
metadata={"help": "An optional input predict data file to do prediction on (a text file)."}, |
|
) |
|
max_target_length: Optional[int] = field( |
|
default=128, |
|
metadata={ |
|
"help": "The maximum total sequence length for target text after tokenization. Sequences longer " |
|
"than this will be truncated, sequences shorter will be padded." |
|
}, |
|
) |
|
val_max_target_length: Optional[int] = field( |
|
default=None, |
|
metadata={ |
|
"help": "The maximum total sequence length for validation target text after tokenization. Sequences longer " |
|
"than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`." |
|
"This argument is also used to override the `max_length` param of `model.generate`, which is used " |
|
"during evaluation." |
|
}, |
|
) |
|
max_train_samples: Optional[int] = field( |
|
default=None, |
|
metadata={ |
|
"help": "For debugging purposes or quicker training, truncate the number of training examples to this " |
|
"value if set." |
|
}, |
|
) |
|
max_eval_samples: Optional[int] = field( |
|
default=None, |
|
metadata={ |
|
"help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this " |
|
"value if set." |
|
}, |
|
) |
|
max_predict_samples: Optional[int] = field( |
|
default=None, |
|
metadata={ |
|
"help": "For debugging purposes or quicker training, truncate the number of prediction examples to this " |
|
"value if set." |
|
}, |
|
) |
|
preprocessing_num_workers: Optional[int] = field( |
|
default=None, |
|
metadata={"help": "The number of processes to use for the preprocessing."}, |
|
) |
|
predict_with_generate: bool = field( |
|
default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."} |
|
) |
|
num_beams: Optional[int] = field( |
|
default=None, |
|
metadata={ |
|
"help": "Number of beams to use for evaluation. This argument will be passed to `model.generate`, " |
|
"which is used during evaluation." |
|
}, |
|
) |
|
overwrite_cache: bool = field( |
|
default=False, metadata={"help": "Overwrite the cached training and evaluation sets"} |
|
) |
|
|
|
def __post_init__(self): |
|
if self.dataset_name is None and self.train_file is None and self.validation_file is None: |
|
raise ValueError("Need either a dataset name or a training/validation file.") |
|
else: |
|
if self.train_file is not None: |
|
extension = self.train_file.split(".")[-1] |
|
assert extension in ["csv", "json"], "`train_file` should be a csv or a json file." |
|
if self.validation_file is not None: |
|
extension = self.validation_file.split(".")[-1] |
|
assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file." |
|
if self.val_max_target_length is None: |
|
self.val_max_target_length = self.max_target_length |
|
|
|
|
|
image_captioning_name_mapping = { |
|
"image_caption_dataset.py": ("image_path", "caption"), |
|
} |
|
|
|
|
|
class TrainState(train_state.TrainState): |
|
dropout_rng: jnp.ndarray |
|
|
|
def replicate(self): |
|
return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng)) |
|
|
|
|
|
def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False): |
|
""" |
|
Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices. |
|
Shuffle batches if `shuffle` is `True`. |
|
""" |
|
steps = len(dataset) // batch_size |
|
|
|
if shuffle: |
|
batch_idx = jax.random.permutation(rng, len(dataset)) |
|
batch_idx = np.asarray(batch_idx) |
|
else: |
|
batch_idx = np.arange(len(dataset)) |
|
|
|
for idx in range(steps): |
|
|
|
start_idx = batch_size * idx |
|
end_idx = batch_size * (idx + 1) |
|
|
|
selected_indices = batch_idx[start_idx:end_idx] |
|
batch = dataset[selected_indices] |
|
batch = shard(batch) |
|
|
|
yield batch |
|
|
|
|
|
def write_metric(summary_writer, metrics, train_time, step, metric_key_prefix="train"): |
|
|
|
if train_time: |
|
summary_writer.scalar("train_time", train_time, step) |
|
|
|
metrics = get_metrics(metrics) |
|
for key, vals in metrics.items(): |
|
tag = f"{metric_key_prefix}_{key}" |
|
for i, val in enumerate(vals): |
|
summary_writer.scalar(tag, val, step - len(vals) + i + 1) |
|
|
|
else: |
|
for metric_name, value in metrics.items(): |
|
summary_writer.scalar(f"{metric_key_prefix}_{metric_name}", value, step) |
|
|
|
|
|
def create_learning_rate_fn( |
|
train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float |
|
) -> Callable[[int], jnp.array]: |
|
"""Returns a linear warmup, linear_decay learning rate function.""" |
|
steps_per_epoch = train_ds_size // train_batch_size |
|
num_train_steps = steps_per_epoch * num_train_epochs |
|
warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps) |
|
decay_fn = optax.linear_schedule( |
|
init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps |
|
) |
|
schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps]) |
|
return schedule_fn |
|
|
|
|
|
def main(): |
|
|
|
|
|
|
|
|
|
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) |
|
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): |
|
|
|
|
|
model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) |
|
else: |
|
model_args, data_args, training_args = parser.parse_args_into_dataclasses() |
|
|
|
if ( |
|
os.path.exists(training_args.output_dir) |
|
and os.listdir(training_args.output_dir) |
|
and training_args.do_train |
|
and not training_args.overwrite_output_dir |
|
): |
|
raise ValueError( |
|
f"Output directory ({training_args.output_dir}) already exists and is not empty." |
|
"Use --overwrite_output_dir to overcome." |
|
) |
|
|
|
|
|
logging.basicConfig( |
|
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
|
datefmt="%m/%d/%Y %H:%M:%S", |
|
level=logging.INFO, |
|
) |
|
|
|
logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR) |
|
if jax.process_index() == 0: |
|
datasets.utils.logging.set_verbosity_warning() |
|
transformers.utils.logging.set_verbosity_info() |
|
else: |
|
datasets.utils.logging.set_verbosity_error() |
|
transformers.utils.logging.set_verbosity_error() |
|
|
|
|
|
logger.info(f"Training/evaluation parameters {training_args}") |
|
|
|
|
|
if training_args.push_to_hub: |
|
if training_args.hub_model_id is None: |
|
repo_name = get_full_repo_name( |
|
Path(training_args.output_dir).absolute().name, token=training_args.hub_token |
|
) |
|
else: |
|
repo_name = training_args.hub_model_id |
|
repo = Repository(training_args.output_dir, clone_from=repo_name) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if data_args.dataset_name is not None: |
|
|
|
dataset = load_dataset( |
|
data_args.dataset_name, |
|
data_args.dataset_config_name, |
|
cache_dir=model_args.cache_dir, |
|
keep_in_memory=False, |
|
data_dir=data_args.data_dir, |
|
) |
|
else: |
|
data_files = {} |
|
if data_args.train_file is not None: |
|
data_files["train"] = data_args.train_file |
|
extension = data_args.train_file.split(".")[-1] |
|
if data_args.validation_file is not None: |
|
data_files["validation"] = data_args.validation_file |
|
extension = data_args.validation_file.split(".")[-1] |
|
if data_args.test_file is not None: |
|
data_files["test"] = data_args.test_file |
|
extension = data_args.test_file.split(".")[-1] |
|
dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir) |
|
|
|
|
|
|
|
|
|
|
|
encoder_cache_dir, decoder_cache_dir = None, None |
|
if model_args.cache_dir: |
|
encoder_cache_dir = os.path.join(model_args.cache_dir, "encoder") |
|
decoder_cache_dir = os.path.join(model_args.cache_dir, "decoder") |
|
|
|
|
|
if model_args.config_name: |
|
config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir) |
|
|
|
elif model_args.model_name_or_path: |
|
config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir) |
|
|
|
else: |
|
|
|
if not model_args.model_type in MODEL_TYPES: |
|
raise ValueError( |
|
f"Unrecognized model identifier: {model_args.model_type}. Should contain one of {', '.join(MODEL_TYPES)}." |
|
) |
|
config_class = CONFIG_MAPPING[model_args.model_type] |
|
|
|
|
|
if hasattr(config_class, "from_encoder_decoder_configs"): |
|
|
|
|
|
if model_args.encoder_config_name: |
|
encoder_config = AutoConfig.from_pretrained(model_args.encoder_config_name, cache_dir=encoder_cache_dir) |
|
|
|
elif model_args.encoder_model_name_or_path: |
|
encoder_config = AutoConfig.from_pretrained(model_args.encoder_model_name_or_path, cache_dir=encoder_cache_dir) |
|
|
|
elif model_args.encoder_model_type: |
|
encoder_config = AutoConfig.for_model(model_args.encoder_model_type) |
|
logger.warning("You are instantiating a new config instance from scratch for the encoder.") |
|
else: |
|
raise ValueError("Encoder Config: if pretrained config or model location is not provided, `encoder_model_type` is required.") |
|
|
|
|
|
if model_args.decoder_config_name: |
|
decoder_config = AutoConfig.from_pretrained(model_args.decoder_config_name, cache_dir=decoder_cache_dir) |
|
|
|
elif model_args.decoder_model_name_or_path: |
|
decoder_config = AutoConfig.from_pretrained(model_args.decoder_model_name_or_path, cache_dir=decoder_cache_dir) |
|
|
|
elif model_args.decoder_model_type: |
|
decoder_config = AutoConfig.for_model(model_args.decoder_model_type) |
|
logger.warning("You are instantiating a new config instance from scratch for the decoder.") |
|
else: |
|
raise ValueError("Decoder Config: if pretrained config or model location is not provided, `decoder_model_type` is required.") |
|
|
|
logger.info("Setting `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config") |
|
decoder_config.is_decoder = True |
|
decoder_config.add_cross_attention = True |
|
|
|
config = config_class.from_encoder_decoder_configs(encoder_config, decoder_config) |
|
|
|
else: |
|
config = config_class() |
|
logger.warning("You are instantiating a new config instance from scratch.") |
|
|
|
decoder_start_token_id = getattr(config, "decoder_start_token_id", None) |
|
if not decoder_start_token_id and getattr(config, "decoder", None): |
|
decoder_start_token_id = getattr(config.decoder, "decoder_start_token_id", None) |
|
bos_token_id = getattr(config, "bos_token_id", None) |
|
if not bos_token_id and getattr(config, "decoder", None): |
|
bos_token_id = getattr(config.decoder, "bos_token_id", None) |
|
eos_token_id = getattr(config, "eos_token_id", None) |
|
if not eos_token_id and getattr(config, "decoder", None): |
|
eos_token_id = getattr(config.decoder, "eos_token_id", None) |
|
pad_token_id = getattr(config, "pad_token_id", None) |
|
if not pad_token_id and getattr(config, "decoder", None): |
|
pad_token_id = getattr(config.decoder, "pad_token_id", None) |
|
|
|
if decoder_start_token_id is None: |
|
decoder_start_token_id = bos_token_id |
|
if pad_token_id is None: |
|
pad_token_id = eos_token_id |
|
|
|
if getattr(config, "decoder", None): |
|
config.decoder.decoder_start_token_id = decoder_start_token_id |
|
config.decoder.bos_token_id = bos_token_id |
|
config.decoder.eos_token_id = eos_token_id |
|
config.decoder.pad_token_id = pad_token_id |
|
|
|
|
|
config.decoder_start_token_id = decoder_start_token_id |
|
config.bos_token_id = bos_token_id |
|
config.eos_token_id = eos_token_id |
|
config.pad_token_id = pad_token_id |
|
|
|
if model_args.model_name_or_path: |
|
model = FlaxAutoModelForVision2Seq.from_pretrained( |
|
model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype) |
|
) |
|
else: |
|
|
|
model = FlaxAutoModelForVision2Seq.from_config(config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)) |
|
model_class = model.__class__ |
|
|
|
|
|
|
|
module = model.module.bind(model.params) |
|
encoder_class_name = type(module.encoder).__name__.replace("Module", "Model") |
|
decoder_class_name = type(module.decoder).__name__.replace("Module", "Model") |
|
encoder_class = getattr(transformers, encoder_class_name, None) |
|
decoder_class = getattr(transformers, decoder_class_name, None) |
|
|
|
if hasattr(model_class, "from_encoder_decoder_pretrained"): |
|
|
|
if model_args.encoder_model_name_or_path: |
|
encoder = encoder_class.from_pretrained( |
|
model_args.encoder_model_name_or_path, |
|
config=config.encoder, |
|
seed=training_args.seed, |
|
dtype=getattr(jnp, model_args.dtype) |
|
) |
|
else: |
|
encoder = encoder_class(config=config.encoder, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)) |
|
logger.warning("You are instantiating a new model instance from scratch for the encoder.") |
|
|
|
if model_args.decoder_model_name_or_path: |
|
decoder = decoder_class.from_pretrained( |
|
model_args.decoder_model_name_or_path, |
|
config=config.decoder, |
|
seed=training_args.seed, |
|
dtype=getattr(jnp, model_args.dtype) |
|
) |
|
else: |
|
decoder = decoder_class(config=config.decoder, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)) |
|
logger.warning("You are instantiating a new model instance from scratch for the decoder.") |
|
|
|
model = model_class.from_encoder_decoder_pretrained( |
|
model_args.encoder_model_name_or_path, |
|
model_args.decoder_model_name_or_path, |
|
encoder_model=encoder, |
|
decoder_model=decoder, |
|
encoder_config=config.encoder, |
|
decoder_config=config.decoder, |
|
encoder_seed=training_args.seed, |
|
decoder_seed=training_args.seed, |
|
encoder_dtype=getattr(jnp, model_args.dtype), |
|
decoder_dtype=getattr(jnp, model_args.dtype), |
|
) |
|
|
|
|
|
model.config.decoder_start_token_id = decoder_start_token_id |
|
model.config.bos_token_id = bos_token_id |
|
model.config.eos_token_id = eos_token_id |
|
model.config.pad_token_id = pad_token_id |
|
|
|
else: |
|
logger.warning("You are instantiating a new model instance from scratch.") |
|
|
|
feature_extractor = None |
|
if model_args.feature_extractor_name: |
|
feature_extractor = AutoFeatureExtractor.from_pretrained( |
|
model_args.feature_extractor_name, cache_dir=model_args.cache_dir, |
|
) |
|
elif model_args.model_name_or_path: |
|
try: |
|
feature_extractor = AutoFeatureExtractor.from_pretrained( |
|
model_args.model_name_or_path, cache_dir=model_args.cache_dir |
|
) |
|
except ValueError as e: |
|
logger.warning(e) |
|
|
|
if not feature_extractor: |
|
if model_args.encoder_model_name_or_path: |
|
feature_extractor = AutoFeatureExtractor.from_pretrained( |
|
model_args.encoder_model_name_or_path, cache_dir=model_args.cache_dir |
|
) |
|
else: |
|
raise ValueError( |
|
"You are instantiating a new feature extractor from scratch. This is not supported by this script." |
|
"You can do it from another script, save it, and load it from here, using --feature_extractor_name." |
|
) |
|
|
|
tokenizer = None |
|
if model_args.tokenizer_name: |
|
tokenizer = AutoTokenizer.from_pretrained( |
|
model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer |
|
) |
|
elif model_args.model_name_or_path: |
|
try: |
|
tokenizer = AutoTokenizer.from_pretrained( |
|
model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer |
|
) |
|
except ValueError as e: |
|
logger.warning(e) |
|
|
|
|
|
if not tokenizer: |
|
if model_args.decoder_model_name_or_path: |
|
tokenizer = AutoTokenizer.from_pretrained( |
|
model_args.decoder_model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer |
|
) |
|
else: |
|
raise ValueError( |
|
"You are instantiating a new tokenizer from scratch. This is not supported by this script." |
|
"You can do it from another script, save it, and load it from here, using --tokenizer_name." |
|
) |
|
tokenizer.pad_token = tokenizer.convert_ids_to_tokens(config.pad_token_id) |
|
|
|
|
|
|
|
if training_args.do_train: |
|
column_names = dataset["train"].column_names |
|
elif training_args.do_eval: |
|
column_names = dataset["validation"].column_names |
|
elif training_args.do_predict: |
|
column_names = dataset["test"].column_names |
|
else: |
|
logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.") |
|
return |
|
|
|
|
|
dataset_columns = image_captioning_name_mapping.get(data_args.dataset_name, None) |
|
if data_args.image_column is None: |
|
assert dataset_columns is not None |
|
image_column = dataset_columns[0] |
|
else: |
|
image_column = data_args.image_column |
|
if image_column not in column_names: |
|
raise ValueError( |
|
f"--image_column' value '{data_args.image_column}' needs to be one of: {', '.join(column_names)}" |
|
) |
|
if data_args.caption_column is None: |
|
assert dataset_columns is not None |
|
caption_column = dataset_columns[1] |
|
else: |
|
caption_column = data_args.caption_column |
|
if caption_column not in column_names: |
|
raise ValueError( |
|
f"--caption_column' value '{data_args.caption_column}' needs to be one of: {', '.join(column_names)}" |
|
) |
|
|
|
|
|
|
|
|
|
model_module = __import__(model.__module__, fromlist=["shift_tokens_right"]) |
|
shift_tokens_right_fn = getattr(model_module, "shift_tokens_right", shift_tokens_right) |
|
|
|
def filter_fn(examples): |
|
"""remove problematic images""" |
|
|
|
bools = [] |
|
for image_file in examples[image_column]: |
|
try: |
|
image = Image.open(image_file) |
|
encoder_inputs = feature_extractor(images=image, return_tensors="np") |
|
bools.append(True) |
|
except: |
|
bools.append(False) |
|
|
|
return bools |
|
|
|
|
|
def tokenization_fn(examples, max_target_length): |
|
"""Run tokenization on captions.""" |
|
|
|
captions = [] |
|
for caption in examples[caption_column]: |
|
captions.append(caption.lower() + ' ' + tokenizer.eos_token) |
|
|
|
targets = captions |
|
|
|
model_inputs = {} |
|
|
|
|
|
with tokenizer.as_target_tokenizer(): |
|
labels = tokenizer( |
|
targets, max_length=max_target_length, padding="max_length", truncation=True, return_tensors="np" |
|
) |
|
|
|
model_inputs["labels"] = labels["input_ids"] |
|
decoder_input_ids = shift_tokens_right_fn( |
|
labels["input_ids"], config.pad_token_id, config.decoder_start_token_id |
|
) |
|
model_inputs["decoder_input_ids"] = np.asarray(decoder_input_ids) |
|
|
|
|
|
model_inputs["decoder_attention_mask"] = labels["attention_mask"] |
|
|
|
model_inputs[image_column] = examples[image_column] |
|
|
|
return model_inputs |
|
|
|
def feature_extraction_fn(examples, check_image=True): |
|
""" |
|
Run feature extraction on images |
|
|
|
If `check_image` is `True`, the examples that fails during `Image.open()` will be caught and discarded. |
|
Otherwise, an exception will be thrown. |
|
""" |
|
|
|
model_inputs = {} |
|
|
|
if check_image: |
|
images = [] |
|
to_keep = [] |
|
for image_file in examples[image_column]: |
|
try: |
|
img = Image.open(image_file) |
|
images.append(img) |
|
to_keep.append(True) |
|
except: |
|
to_keep.append(False) |
|
|
|
for k, v in examples.items(): |
|
if k != image_column: |
|
model_inputs[k] = v[to_keep] |
|
else: |
|
images = [Image.open(image_file) for image_file in examples[image_column]] |
|
|
|
encoder_inputs = feature_extractor(images=images, return_tensors="np") |
|
model_inputs["pixel_values"] = encoder_inputs.pixel_values |
|
|
|
return model_inputs |
|
|
|
def preprocess_fn(examples, max_target_length, check_image=True): |
|
"""Run tokenization + image feature extraction""" |
|
|
|
model_inputs = {} |
|
|
|
model_inputs.update(tokenization_fn(examples, max_target_length)) |
|
model_inputs.update(feature_extraction_fn(model_inputs, check_image=check_image)) |
|
|
|
model_inputs.pop(image_column) |
|
|
|
return model_inputs |
|
|
|
features = datasets.Features( |
|
{ |
|
"pixel_values": datasets.Array3D( |
|
shape=( |
|
getattr(config.encoder, "num_channels", 3), |
|
config.encoder.image_size, |
|
config.encoder.image_size, |
|
), |
|
dtype="float32", |
|
), |
|
"labels": datasets.Sequence(feature=datasets.Value(dtype='int32', id=None), length=-1, id=None), |
|
"decoder_input_ids": datasets.Sequence(feature=datasets.Value(dtype='int32', id=None), length=-1, id=None), |
|
"decoder_attention_mask": datasets.Sequence(feature=datasets.Value(dtype='int32', id=None), length=-1, id=None), |
|
} |
|
) |
|
|
|
|
|
run_feat_ext_at_beginning = training_args.block_size == 0 |
|
|
|
function_kwarg = preprocess_fn if run_feat_ext_at_beginning else tokenization_fn |
|
|
|
features_kwarg = features if run_feat_ext_at_beginning else None |
|
|
|
remove_columns_kwarg = [x for x in column_names if x != image_column or run_feat_ext_at_beginning] |
|
processor_names = "tokenizer and feature extractor" if run_feat_ext_at_beginning else "tokenizer" |
|
|
|
if training_args.do_train: |
|
if "train" not in dataset: |
|
raise ValueError("--do_train requires a train dataset") |
|
train_dataset = dataset["train"] |
|
if data_args.max_train_samples is not None: |
|
train_dataset = train_dataset.select(range(data_args.max_train_samples)) |
|
|
|
|
|
|
|
if not run_feat_ext_at_beginning: |
|
train_dataset = train_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers) |
|
train_dataset = train_dataset.map( |
|
function=function_kwarg, |
|
batched=True, |
|
num_proc=data_args.preprocessing_num_workers, |
|
|
|
remove_columns=remove_columns_kwarg, |
|
load_from_cache_file=not data_args.overwrite_cache, |
|
desc=f"Running {processor_names} on train dataset", |
|
fn_kwargs={"max_target_length": data_args.max_target_length}, |
|
features=features_kwarg, |
|
) |
|
if run_feat_ext_at_beginning: |
|
|
|
train_dataset = train_dataset.with_format("numpy") |
|
|
|
if training_args.do_eval: |
|
if "validation" not in dataset: |
|
raise ValueError("--do_eval requires a validation dataset") |
|
eval_dataset = dataset["validation"] |
|
if data_args.max_eval_samples is not None: |
|
eval_dataset = eval_dataset.select(range(data_args.max_eval_samples)) |
|
|
|
|
|
|
|
if not run_feat_ext_at_beginning: |
|
eval_dataset = eval_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers) |
|
eval_dataset = eval_dataset.map( |
|
function=function_kwarg, |
|
batched=True, |
|
num_proc=data_args.preprocessing_num_workers, |
|
|
|
remove_columns=remove_columns_kwarg, |
|
load_from_cache_file=not data_args.overwrite_cache, |
|
desc=f"Running {processor_names} on validation dataset", |
|
fn_kwargs={"max_target_length": data_args.val_max_target_length}, |
|
features=features_kwarg, |
|
) |
|
if run_feat_ext_at_beginning: |
|
|
|
eval_dataset = eval_dataset.with_format("numpy") |
|
|
|
if training_args.do_predict: |
|
if "test" not in dataset: |
|
raise ValueError("--do_predict requires a test dataset") |
|
predict_dataset = dataset["test"] |
|
if data_args.max_predict_samples is not None: |
|
predict_dataset = predict_dataset.select(range(data_args.max_predict_samples)) |
|
|
|
|
|
|
|
if not run_feat_ext_at_beginning: |
|
predict_dataset = predict_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers) |
|
predict_dataset = predict_dataset.map( |
|
function=function_kwarg, |
|
batched=True, |
|
num_proc=data_args.preprocessing_num_workers, |
|
|
|
remove_columns=remove_columns_kwarg, |
|
load_from_cache_file=not data_args.overwrite_cache, |
|
desc=f"Running {processor_names} on prediction dataset", |
|
fn_kwargs={"max_target_length": data_args.val_max_target_length}, |
|
features=features_kwarg, |
|
) |
|
if run_feat_ext_at_beginning: |
|
|
|
predict_dataset = predict_dataset.with_format("numpy") |
|
|
|
|
|
|
|
train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count() |
|
|
|
if training_args.block_size % train_batch_size > 0: |
|
raise ValueError(f"`training_args.block_size` needs to be a multiple of the global batch size. Got {training_args.block_size} and {train_batch_size} instead.") |
|
|
|
if training_args.do_train: |
|
steps_per_epoch = len(train_dataset) // train_batch_size |
|
num_train_examples_per_epoch = steps_per_epoch * train_batch_size |
|
num_epochs = int(training_args.num_train_epochs) |
|
total_train_steps = steps_per_epoch * num_epochs |
|
else: |
|
num_train_examples_per_epoch = 0 |
|
|
|
eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count() |
|
|
|
if training_args.do_eval: |
|
num_eval_examples = len(eval_dataset) |
|
eval_steps = num_eval_examples // eval_batch_size |
|
|
|
if training_args.do_predict: |
|
num_test_examples = len(predict_dataset) |
|
test_steps = num_test_examples // eval_batch_size |
|
|
|
def blockwise_data_loader( |
|
rng: jax.random.PRNGKey, |
|
ds: Dataset, |
|
block_size: int, |
|
batch_size: int, |
|
shuffle: bool = False, |
|
keep_in_memory: bool = False, |
|
split: str = "" |
|
): |
|
""" |
|
Wrap the simple `data_loader` in a block-wise way if `block_size` > 0, else it's the same as `data_loader`. |
|
|
|
If `block_size` > 0, it requires `ds` to have a column that gives image paths in order to perform image feature |
|
extraction (with the column name being specified by `image_column`). The tokenization should be done before |
|
training in this case. |
|
""" |
|
|
|
if shuffle: |
|
indices = jax.random.permutation(rng, len(ds)) |
|
indices = np.asarray(indices) |
|
else: |
|
indices = np.arange(len(ds)) |
|
|
|
_block_size = len(ds) if not block_size else block_size |
|
|
|
steps_per_block = _block_size // batch_size |
|
num_examples = len(ds) |
|
steps = num_examples // batch_size |
|
num_splits = steps // steps_per_block + int(steps % steps_per_block > 0) |
|
|
|
for idx in range(num_splits): |
|
|
|
if not block_size: |
|
_ds = ds |
|
else: |
|
|
|
start_idx = block_size * idx |
|
end_idx = block_size * (idx + 1) |
|
|
|
selected_indices = indices[start_idx:end_idx] |
|
|
|
_ds = ds.select(selected_indices) |
|
|
|
_ds = _ds.map( |
|
feature_extraction_fn, |
|
batched=True, |
|
num_proc=data_args.preprocessing_num_workers, |
|
remove_columns=[image_column], |
|
load_from_cache_file=not data_args.overwrite_cache, |
|
features=features, |
|
keep_in_memory=keep_in_memory, |
|
|
|
fn_kwargs={"check_image": False}, |
|
desc=f"Running feature extraction on {split} dataset".replace(" ", " "), |
|
) |
|
_ds = _ds.with_format("numpy") |
|
|
|
|
|
loader = data_loader(rng, _ds, batch_size=batch_size, shuffle=False) |
|
|
|
for batch in loader: |
|
yield batch |
|
|
|
|
|
metric = load_metric("rouge") |
|
|
|
def postprocess_text(preds, labels): |
|
preds = [pred.strip() for pred in preds] |
|
labels = [label.strip() for label in labels] |
|
|
|
|
|
preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds] |
|
labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels] |
|
|
|
return preds, labels |
|
|
|
def compute_metrics(preds, labels): |
|
decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True) |
|
decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) |
|
|
|
|
|
decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels) |
|
|
|
result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True) |
|
|
|
result = {key: value.mid.fmeasure * 100 for key, value in result.items()} |
|
|
|
prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds] |
|
result["gen_len"] = np.mean(prediction_lens) |
|
result = {k: round(v, 6) for k, v in result.items()} |
|
|
|
return result, decoded_preds, decoded_labels |
|
|
|
|
|
has_tensorboard = is_tensorboard_available() |
|
if has_tensorboard and jax.process_index() == 0: |
|
try: |
|
from flax.metrics.tensorboard import SummaryWriter |
|
|
|
summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir)) |
|
except ImportError as ie: |
|
has_tensorboard = False |
|
logger.warning( |
|
f"Unable to display metrics through TensorBoard because some package are not installed: {ie}" |
|
) |
|
else: |
|
logger.warning( |
|
"Unable to display metrics through TensorBoard because the package is not installed: " |
|
"Please run pip install tensorboard to enable." |
|
) |
|
|
|
|
|
rng = jax.random.PRNGKey(training_args.seed) |
|
rng, dropout_rng = jax.random.split(rng) |
|
|
|
|
|
linear_decay_lr_schedule_fn = create_learning_rate_fn( |
|
num_train_examples_per_epoch, |
|
train_batch_size, |
|
training_args.num_train_epochs, |
|
training_args.warmup_steps, |
|
training_args.learning_rate, |
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def decay_mask_fn(params): |
|
flat_params = traverse_util.flatten_dict(params) |
|
layer_norm_params = [ |
|
(name, "scale") for name in ["self_attn_layer_norm", "layernorm_embedding", "final_layer_norm"] |
|
] |
|
flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_params) for path in flat_params} |
|
return traverse_util.unflatten_dict(flat_mask) |
|
|
|
|
|
adamw = optax.adamw( |
|
learning_rate=linear_decay_lr_schedule_fn, |
|
b1=training_args.adam_beta1, |
|
b2=training_args.adam_beta2, |
|
eps=training_args.adam_epsilon, |
|
weight_decay=training_args.weight_decay, |
|
mask=decay_mask_fn, |
|
) |
|
|
|
|
|
state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng) |
|
|
|
|
|
def loss_fn(logits, labels, padding_mask, label_smoothing_factor=0.0): |
|
""" |
|
The label smoothing implementation is adapted from Flax's official example: |
|
https://github.com/google/flax/blob/87a211135c6a377c8f29048a1cac3840e38b9da4/examples/wmt/train.py#L104 |
|
""" |
|
vocab_size = logits.shape[-1] |
|
confidence = 1.0 - label_smoothing_factor |
|
low_confidence = (1.0 - confidence) / (vocab_size - 1) |
|
normalizing_constant = -( |
|
confidence * jnp.log(confidence) + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20) |
|
) |
|
soft_labels = onehot(labels, vocab_size, on_value=confidence, off_value=low_confidence) |
|
|
|
loss = optax.softmax_cross_entropy(logits, soft_labels) |
|
loss = loss - normalizing_constant |
|
|
|
|
|
loss = loss * padding_mask |
|
loss = loss.sum() / padding_mask.sum() |
|
return loss |
|
|
|
|
|
def train_step(state, batch, label_smoothing_factor=0.0): |
|
dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng) |
|
|
|
def compute_loss(params): |
|
labels = batch.pop("labels") |
|
logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0] |
|
loss = loss_fn(logits, labels, batch["decoder_attention_mask"], label_smoothing_factor) |
|
return loss |
|
|
|
grad_fn = jax.value_and_grad(compute_loss) |
|
loss, grad = grad_fn(state.params) |
|
grad = jax.lax.pmean(grad, "batch") |
|
|
|
new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng) |
|
|
|
metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)} |
|
metrics = jax.lax.pmean(metrics, axis_name="batch") |
|
|
|
return new_state, metrics |
|
|
|
|
|
def eval_step(params, batch, label_smoothing_factor=0.0): |
|
labels = batch.pop("labels") |
|
logits = model(**batch, params=params, train=False)[0] |
|
loss = loss_fn(logits, labels, batch["decoder_attention_mask"], label_smoothing_factor) |
|
|
|
|
|
metrics = {"loss": loss} |
|
metrics = jax.lax.pmean(metrics, axis_name="batch") |
|
return metrics |
|
|
|
|
|
max_length = ( |
|
data_args.val_max_target_length if data_args.val_max_target_length is not None else model.config.max_length |
|
) |
|
num_beams = data_args.num_beams if data_args.num_beams is not None else model.config.num_beams |
|
gen_kwargs = {"max_length": max_length, "num_beams": num_beams} |
|
|
|
def generate_step(params, batch): |
|
model.params = params |
|
output_ids = model.generate(batch['pixel_values'], **gen_kwargs) |
|
return output_ids.sequences |
|
|
|
|
|
p_train_step = jax.pmap( |
|
partial(train_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch", donate_argnums=(0,) |
|
) |
|
p_eval_step = jax.pmap(partial(eval_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch") |
|
p_generate_step = jax.pmap(generate_step, "batch") |
|
|
|
|
|
state = state.replicate() |
|
|
|
if training_args.do_train: |
|
logger.info("***** Running training *****") |
|
logger.info(f" Num train examples = {num_train_examples_per_epoch}") |
|
logger.info(f" Num Epochs = {num_epochs}") |
|
logger.info(f" Instantaneous train batch size per device = {training_args.per_device_train_batch_size}") |
|
logger.info(f" Total train batch size (w. parallel & distributed) = {train_batch_size}") |
|
logger.info(f" Optimization steps per epoch = {steps_per_epoch}") |
|
logger.info(f" Total optimization steps = {total_train_steps}") |
|
if training_args.do_eval: |
|
logger.info(f" Num evaluation examples = {num_eval_examples}") |
|
logger.info(f" Instantaneous evaluation batch size per device = {training_args.per_device_eval_batch_size}") |
|
logger.info(f" Total evaluation batch size (w. parallel & distributed) = {eval_batch_size}") |
|
logger.info(f" Evaluation steps = {eval_steps}") |
|
if training_args.do_predict: |
|
logger.info(f" Num test examples = {num_test_examples}") |
|
logger.info(f" Instantaneous test batch size per device = {training_args.per_device_eval_batch_size}") |
|
logger.info(f" Total test batch size (w. parallel & distributed) = {eval_batch_size}") |
|
logger.info(f" Test steps = {test_steps}") |
|
|
|
|
|
if not os.path.isdir(os.path.join(training_args.output_dir)): |
|
os.makedirs(os.path.join(training_args.output_dir), exist_ok=True) |
|
|
|
def save_ckpt(ckpt_dir: str, commit_msg: str = ""): |
|
"""save checkpoints and push to Hugging Face Hub if specified""" |
|
|
|
|
|
if jax.process_index() == 0: |
|
params = jax.device_get(jax.tree_map(lambda x: x[0], state.params)) |
|
model.save_pretrained(os.path.join(training_args.output_dir, ckpt_dir), params=params) |
|
tokenizer.save_pretrained(os.path.join(training_args.output_dir, ckpt_dir)) |
|
if training_args.push_to_hub: |
|
repo.push_to_hub(commit_message=commit_msg, blocking=False) |
|
|
|
def evaluation_loop(rng: jax.random.PRNGKey, dataset: Dataset, metric_key_prefix: str = "eval", ckpt_dir: str = "", is_prediction=False): |
|
|
|
logger.info(f"*** {'Predict' if is_prediction else 'Evaluate'} ***") |
|
|
|
metrics = [] |
|
preds = [] |
|
labels = [] |
|
|
|
batches = blockwise_data_loader( |
|
rng, |
|
dataset, |
|
block_size=training_args.block_size, |
|
batch_size=eval_batch_size, |
|
keep_in_memory=False, |
|
shuffle=False, |
|
split="prediction" if is_prediction else "validation", |
|
) |
|
steps = len(dataset) // eval_batch_size |
|
for _ in tqdm(range(steps), desc=f"{'Predicting' if is_prediction else 'Evaluating'}...", position=2, leave=False): |
|
|
|
batch = next(batches) |
|
_labels = batch.get("labels", None) |
|
if not is_prediction and _labels is None: |
|
raise ValueError("Evaluation requires the validation dataset to have `labels`") |
|
|
|
if _labels is not None: |
|
_metrics = p_eval_step(state.params, batch) |
|
metrics.append(_metrics) |
|
|
|
|
|
if data_args.predict_with_generate: |
|
generated_ids = p_generate_step(state.params, batch) |
|
preds.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"]))) |
|
if _labels is not None: |
|
labels.extend(jax.device_get(_labels.reshape(-1, _labels.shape[-1]))) |
|
|
|
if metrics: |
|
|
|
metrics = get_metrics(metrics) |
|
metrics = jax.tree_map(jnp.mean, metrics) |
|
|
|
|
|
generations = [] |
|
rouge_desc = "" |
|
if data_args.predict_with_generate: |
|
if labels: |
|
rouge_metrics, decoded_preds, decoded_labels = compute_metrics(preds, labels) |
|
metrics.update(rouge_metrics) |
|
rouge_desc = " ".join([f"{'Predict' if is_prediction else 'Eval'} {key}: {value} |" for key, value in rouge_metrics.items()]) |
|
for pred, label in zip(decoded_preds, decoded_labels): |
|
pred = pred.replace("\n", " ") |
|
label = label.replace("\n", " ") |
|
generations.append({"label": label, "pred": pred}) |
|
else: |
|
decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True) |
|
|
|
decoded_preds = [pred.strip() for pred in decoded_preds] |
|
|
|
decoded_preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in decoded_preds] |
|
for pred in decoded_preds: |
|
pred = pred.replace("\n", " ") |
|
generations.append({"pred": pred}) |
|
|
|
if metrics: |
|
|
|
desc = f"{'Predict' if is_prediction else 'Eval'} Loss: {metrics['loss']} | {rouge_desc})" |
|
if training_args.do_train and not is_prediction: |
|
desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | " + desc |
|
epochs.write(desc) |
|
epochs.desc = desc |
|
logger.info(desc) |
|
|
|
if jax.process_index() == 0: |
|
|
|
if not os.path.isdir(os.path.join(training_args.output_dir, ckpt_dir)): |
|
os.makedirs(os.path.join(training_args.output_dir, ckpt_dir), exist_ok=True) |
|
|
|
if metrics: |
|
|
|
|
|
if has_tensorboard and training_args.do_train: |
|
write_metric(summary_writer, metrics, train_time=None, step=cur_step, metric_key_prefix=metric_key_prefix) |
|
|
|
|
|
metrics = {f"{metric_key_prefix}_{metric_name}": round(value.item(), 6) for metric_name, value in metrics.items()} |
|
_path = os.path.join(training_args.output_dir, ckpt_dir, f"{metric_key_prefix}_results.json") |
|
with open(_path, "w") as f: |
|
json.dump(metrics, f, indent=4, sort_keys=True) |
|
|
|
|
|
with open(os.path.join(training_args.output_dir, 'log'), 'a', encoding='UTF-8') as fp: |
|
fp.write(desc + '\n') |
|
|
|
|
|
if generations: |
|
with open(os.path.join(training_args.output_dir, ckpt_dir, f'{metric_key_prefix}_generation.json'), 'w', encoding='UTF-8') as fp: |
|
json.dump(generations, fp, ensure_ascii=False, indent=4) |
|
|
|
def evaluate(rng: jax.random.PRNGKey, dataset: Dataset, ckpt_dir: str = ""): |
|
evaluation_loop(rng, dataset, metric_key_prefix='eval', ckpt_dir=ckpt_dir) |
|
|
|
def predict(rng: jax.random.PRNGKey, dataset: Dataset): |
|
evaluation_loop(rng, dataset, metric_key_prefix='test', is_prediction=True) |
|
|
|
input_rng = None |
|
|
|
if training_args.do_train: |
|
|
|
cur_step = 0 |
|
train_time = 0 |
|
epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0) |
|
|
|
for epoch in epochs: |
|
|
|
|
|
|
|
|
|
rng, input_rng = jax.random.split(rng) |
|
|
|
train_metrics = [] |
|
|
|
train_batches = blockwise_data_loader( |
|
input_rng, |
|
train_dataset, |
|
block_size=training_args.block_size, |
|
batch_size=train_batch_size, |
|
keep_in_memory=True, |
|
shuffle=True, |
|
split="train" |
|
) |
|
|
|
|
|
for (batch_idx, _) in enumerate(tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False)): |
|
|
|
cur_step += 1 |
|
batch = next(train_batches) |
|
batch_start = time.time() |
|
state, train_metric = p_train_step(state, batch) |
|
train_metrics.append(train_metric) |
|
train_time += time.time() - batch_start |
|
|
|
time_per_step = train_time / cur_step |
|
_train_metric = unreplicate(train_metric) |
|
desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | Loss: {_train_metric['loss']} | Learning Rate: {_train_metric['learning_rate']} | Time per step: {time_per_step})" |
|
epochs.desc = desc |
|
epochs.write(desc) |
|
|
|
|
|
if training_args.logging_steps > 0 and cur_step % training_args.logging_steps == 0: |
|
|
|
logger.info(desc) |
|
|
|
with open(os.path.join(training_args.output_dir, 'log'), 'a', encoding='UTF-8') as fp: |
|
fp.write(desc + '\n') |
|
|
|
|
|
if has_tensorboard and jax.process_index() == 0: |
|
write_metric(summary_writer, train_metrics, train_time=train_time, step=cur_step, metric_key_prefix="train") |
|
|
|
|
|
|
|
if training_args.do_eval and (training_args.eval_steps is not None and training_args.eval_steps > 0) and cur_step % training_args.eval_steps == 0: |
|
ckpt_dir = f"ckpt_epoch_{epoch + 1}_step_{cur_step}" |
|
commit_msg = f"Saving weights and logs of epoch {epoch + 1} - step {cur_step}" |
|
evaluate(input_rng, eval_dataset, ckpt_dir) |
|
save_ckpt(ckpt_dir=ckpt_dir, commit_msg=commit_msg) |
|
|
|
|
|
|
|
|
|
if training_args.logging_steps <= 0: |
|
|
|
logger.info(desc) |
|
|
|
with open(os.path.join(training_args.output_dir, 'log'), 'a', encoding='UTF-8') as fp: |
|
fp.write(desc + '\n') |
|
|
|
|
|
if has_tensorboard and jax.process_index() == 0: |
|
write_metric(summary_writer, train_metrics, train_time=train_time, step=cur_step, metric_key_prefix="train") |
|
|
|
|
|
|
|
if training_args.do_eval and (training_args.eval_steps is None or training_args.eval_steps <= 0): |
|
ckpt_dir = f"ckpt_epoch_{epoch + 1}_step_{cur_step}" |
|
commit_msg = f"Saving weights and logs of epoch {epoch + 1} - step {cur_step}" |
|
evaluate(input_rng, eval_dataset, ckpt_dir) |
|
save_ckpt(ckpt_dir=ckpt_dir, commit_msg=commit_msg) |
|
|
|
|
|
|
|
|
|
if input_rng is None: |
|
rng, input_rng = jax.random.split(rng) |
|
|
|
|
|
if training_args.do_eval and not training_args.do_train: |
|
evaluate(input_rng, eval_dataset) |
|
|
|
|
|
if training_args.do_predict: |
|
predict(input_rng, predict_dataset) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
main() |
|
|