glm

libai/optim/__init__.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .build import build_optimizer, get_default_optimizer_params

libai/optim/build.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from collections import defaultdict
+from typing import Any, Dict, List
+
+import oneflow as flow
+
+from libai.config import instantiate
+from libai.layers import LayerNorm
+
+# --------------------------------------------------------
+# References:
+# https://github.com/facebookresearch/detectron2/blob/main/detectron2/solver/build.py
+# --------------------------------------------------------
+
+
+def build_optimizer(cfg, model):
+    """
+    Build an optimizer from config.
+    """
+
+    cfg.params.model = model
+    optim = instantiate(cfg)
+    return optim
+
+
+def get_default_optimizer_params(
+    model,
+    base_lr=None,
+    weight_decay=None,
+    weight_decay_norm=None,
+    weight_decay_bias=None,
+    clip_grad_max_norm=None,
+    clip_grad_norm_type=None,
+    overrides=None,
+):
+    """
+    Get default param list for optimizer, with suport for a few types of overrides.
+    If no overrides are needed, it is equivalent to `model.parameters()`.
+
+    Arguments:
+        base_lr: lr for every group by default. Can be omitted to use the one in optimizer.
+        weight_decay: weight decay for every group by default. Can be omitted to use the one
+            in optimizer.
+        weight_decay_norm: override weight decay for params in normalization layers
+        weight_decay_bias: override weight decay for bias parameters
+        overrides: if not `None`, provides values for optimizer hyperparameters
+            (LR, weight decay) for module parameters with a given name; e.g.
+            ``{"embedding": {"lr": 0.01, "weight_decay": 0.1}}`` will set the LR and
+            weight decay values for all module parameters named `embedding`.
+
+    For common transformer models, ``weight_decay_norm`` and ``weight_decay_bias``
+    are usually set to 0.
+
+    Example:
+    ::
+
+        flow.optim.AdamW(
+            get_default_optimizer_params(model, weight_decay_norm=0, weight_decay_bias=0),
+            lr=0.01,
+            weight_decay=1e-4
+        )
+    """
+    if overrides is None:
+        overrides = {}
+    defaults = {}
+    if base_lr is not None:
+        defaults["lr"] = base_lr
+    if weight_decay is not None:
+        defaults["weight_decay"] = weight_decay
+    if clip_grad_max_norm is not None and clip_grad_norm_type is not None:
+        defaults["clip_grad_max_norm"] = clip_grad_max_norm
+        defaults["clip_grad_norm_type"] = clip_grad_norm_type
+    bias_overrides = {}
+    if weight_decay_bias is not None:
+        bias_overrides["weight_decay"] = weight_decay_bias
+    if len(bias_overrides):
+        if "bias" in overrides:
+            raise ValueError("Conflicting overrides for 'bias'")
+        overrides["bias"] = bias_overrides
+
+    norm_module_types = (
+        LayerNorm,
+        flow.nn.BatchNorm1d,
+        flow.nn.BatchNorm2d,
+        flow.nn.BatchNorm3d,
+        flow.nn.GroupNorm,
+        flow.nn.InstanceNorm1d,
+        flow.nn.InstanceNorm2d,
+        flow.nn.InstanceNorm3d,
+        flow.nn.FusedBatchNorm1d,
+        flow.nn.FusedBatchNorm2d,
+        flow.nn.FusedBatchNorm3d,
+    )
+    params = []
+    memo = set()
+    for module in model.modules():
+        for model_param_name, value in module.named_parameters(recurse=False):
+            if not value.requires_grad:
+                continue
+            # Avoid duplicating parameters
+            if value in memo:
+                continue
+            memo.add(value)
+
+            hyperparams = copy.copy(defaults)
+            if isinstance(module, norm_module_types) and weight_decay_norm is not None:
+                hyperparams["weight_decay"] = weight_decay_norm
+            hyperparams.update(overrides.get(model_param_name, {}))
+            params.append({"params": [value], **hyperparams})
+    return reduce_param_groups(params)
+
+
+def _expand_param_groups(params: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Transform parameter groups into per-parameter structure.
+    Later items in `params` can overwrite parameters set in previous items.
+    """
+    ret = defaultdict(dict)
+    for item in params:
+        assert "params" in item
+        cur_params = {x: y for x, y in item.items() if x != "params"}
+        for param in item["params"]:
+            ret[param].update({"params": [param], **cur_params})
+    return list(ret.values())
+
+
+def reduce_param_groups(params: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Reorganize the parameter groups and merge duplicated groups.
+    The number of parameter groups needs to be as small as possible in order
+    to efficiently use the OneFlow multi-tensor optimizer. Therefore instead
+    of using a parameter_group per single parameter, we reorganize the
+    parameter groups and merge duplicated groups. This approach speeds
+    up multi-tensor optimizer significantly.
+    """
+    params = _expand_param_groups(params)
+    groups = defaultdict(list)  # re-group all parameter groups by their hyperparams
+    for item in params:
+        cur_params = tuple((x, y) for x, y in item.items() if x != "params")
+        groups[cur_params].extend(item["params"])
+    ret = []
+    for param_keys, param_values in groups.items():
+        cur = {kv[0]: kv[1] for kv in param_keys}
+        cur["params"] = param_values
+        ret.append(cur)
+    return ret

libai/scheduler/__init__.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .build import build_lr_scheduler
+from .lr_scheduler import (
+    WarmupCosineAnnealingLR,
+    WarmupCosineLR,
+    WarmupExponentialLR,
+    WarmupMultiStepLR,
+    WarmupPolynomialLR,
+    WarmupStepLR,
+)

libai/scheduler/build.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from libai.config import instantiate
+
+
+def build_lr_scheduler(cfg, optimizer):
+    """Build learning rate scheduler, defined by ``cfg``."""
+    cfg.optimizer = optimizer
+    scheduler = instantiate(cfg)
+    return scheduler

libai/scheduler/lr_scheduler.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import oneflow as flow
+
+logger = logging.getLogger(__name__)
+
+
+def WarmupCosineLR(
+    optimizer: flow.optim.Optimizer,
+    max_iter: int,
+    warmup_factor: float,
+    warmup_iter: int,
+    alpha: float = 0.0,
+    warmup_method: str = "linear",
+):
+    """Create a schedule with a learning rate that decreases following
+    the values of the Cosine function between the initial lr set in the
+    optimizer to 0, after a warmup period during which it increases linearly
+    between 0 and the initial lr set in the optimizer.
+
+    Args:
+        optimizer (flow.optim.Optimizer): Wrapped optimizer.
+        max_iter (int): Total training iters.
+        warmup_factor (float): The warmup factor.
+        warmup_iter (int): The number of warmup steps.
+        alpha (float, optional): The learning rate scale factor (:math:`\\alpha`). Defaults to 0.0.
+        warmup_method (str, optional): The method of warmup, you can choose "linear" or "constant".
+            In linear mode, the multiplication factor starts with warmup_factor in
+            the first epoch and then inreases linearly to reach 1. Defaults to "linear".
+    """
+    cosine_decay_lr = flow.optim.lr_scheduler.CosineDecayLR(
+        optimizer, decay_steps=max_iter, alpha=alpha
+    )
+    if warmup_iter == 0:
+        logger.warning("warmup iters equals to zero, return CosineLR")
+        return cosine_decay_lr
+    elif warmup_iter > max_iter:
+        logger.warning("warmup iters is larger than the total training iters")
+    warmup_cosine_lr = flow.optim.lr_scheduler.WarmUpLR(
+        cosine_decay_lr,
+        warmup_factor=warmup_factor,
+        warmup_iters=warmup_iter,
+        warmup_method=warmup_method,
+    )
+    return warmup_cosine_lr
+
+
+def WarmupCosineAnnealingLR(
+    optimizer: flow.optim.Optimizer,
+    max_iter: int,
+    warmup_factor: float,
+    warmup_iter: int,
+    eta_min: float = 0.0,
+    warmup_method: str = "linear",
+):
+    """Create a schedule with a learning rate that decreases following
+    the values of the Cosine Annealing function between the initial
+    lr set in the optimizer to 0, after a warmup period during which
+    it increases linearly between 0 and the initial lr set in the optimizer.
+
+    Args:
+        optimizer (flow.optim.Optimizer): Wrapped optimizer.
+        max_iter (int): Total training iters.
+        warmup_factor (float): The warmup factor.
+        warmup_iter (int): The number of warmup steps.
+        eta_min (float, optional): Minimum learning rate. Defaults to 0.0.
+        warmup_method (str, optional): The method of warmup, you can choose "linear" or "constant".
+            In linear mode, the multiplication factor starts with warmup_factor in the first epoch
+            and then inreases linearly to reach 1. Defaults to "linear".
+    """
+    cosine_annealing_lr = flow.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=max_iter, eta_min=eta_min
+    )
+    if warmup_iter == 0:
+        logger.warning("warmup iters equals to zero, return CosineAnnealingLR")
+        return cosine_annealing_lr
+    warmup_cosine_annealing_lr = flow.optim.lr_scheduler.WarmUpLR(
+        cosine_annealing_lr,
+        warmup_factor=warmup_factor,
+        warmup_iters=warmup_iter,
+        warmup_method=warmup_method,
+    )
+    return warmup_cosine_annealing_lr
+
+
+def WarmupStepLR(
+    optimizer: flow.optim.Optimizer,
+    max_iter: int,
+    warmup_factor: float,
+    warmup_iter: int,
+    step_size: int,
+    gamma: float = 0.1,
+    warmup_method: str = "linear",
+):
+    """Create a schedule with a learning rate that decreases following the values of the Step
+    function between the initial lr set in the optimizer to 0, after a warmup period during which
+    it increases linearly between 0 and the initial lr set in the optimizer.
+    Args:
+        optimizer (flow.optim.Optimizer): Wrapped optimizer.
+        max_iter (int): Total training iters.
+        warmup_factor (float): The warmup factor.
+        warmup_iter (int): The number of warmup steps.
+        step_size (int): Period of learning rate decay.
+        gamma (float, optional): Multiplicative factor of learning rate decay. Defaults to 0.1.
+        warmup_method (str, optional): The method of warmup, you can choose "linear" or "constant".
+            In linear mode, the multiplication factor starts with warmup_factor in the first
+            epoch and then inreases linearly to reach 1. Defaults to "linear".
+    """
+    step_lr = flow.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)
+    if warmup_iter == 0:
+        logger.warning("warmup iters equals to zero, return StepLR")
+        return step_lr
+    warmup_step_lr = flow.optim.lr_scheduler.WarmUpLR(
+        step_lr,
+        warmup_factor=warmup_factor,
+        warmup_iters=warmup_iter,
+        warmup_method=warmup_method,
+    )
+    return warmup_step_lr
+
+
+def WarmupMultiStepLR(
+    optimizer: flow.optim.Optimizer,
+    max_iter: int,
+    warmup_factor: float,
+    warmup_iter: int,
+    milestones: list,
+    gamma: float = 0.1,
+    warmup_method: str = "linear",
+):
+    """Create a schedule with a learning rate that decreases following the values of the MultiStep
+    function between the initial lr set in the optimizer to 0, after a warmup period during which
+    it increases linearly between 0 and the initial lr set in the optimizer.
+
+    Args:
+        optimizer (flow.optim.Optimizer): Wrapped optimizer.
+        max_iter (int): Total training iters.
+        warmup_factor (float): The warmup factor.
+        warmup_iter (int): The number of warmup steps.
+        milestones (list): List of step indices. Must be increasing.
+        gamma (float, optional): Multiplicative factor of learning rate decay. Defaults to 0.1.
+        warmup_method (str, optional): The method of warmup, you can choose "linear" or "constant".
+            In linear mode, the multiplication factor starts with warmup_factor in the first
+            epoch and then inreases linearly to reach 1. Defaults to "linear".
+    """
+    multistep_lr = flow.optim.lr_scheduler.MultiStepLR(
+        optimizer, milestones=milestones, gamma=gamma
+    )
+    if warmup_iter == 0:
+        logger.warning("warmup iters equals to zero, return MultiStepLR")
+        return multistep_lr
+    warmup_multistep_lr = flow.optim.lr_scheduler.WarmUpLR(
+        multistep_lr,
+        warmup_factor=warmup_factor,
+        warmup_iters=warmup_iter,
+        warmup_method=warmup_method,
+    )
+    return warmup_multistep_lr
+
+
+def WarmupExponentialLR(
+    optimizer: flow.optim.Optimizer,
+    max_iter: int,
+    gamma: float,
+    warmup_factor: float,
+    warmup_iter: int,
+    warmup_method: str = "linear",
+):
+    """Create a schedule with a learning rate that decreases following the values of
+    the Exponential function between the initial lr set in the optimizer to 0,
+    after a warmup period during which it increases linearly between 0 and the
+    initial lr set in the optimizer.
+
+    Args:
+        optimizer (flow.optim.Optimizer): Wrapped optimizer.
+        max_iter (int): Total training iters.
+        gamma (float): Multiplicative factor of learning rate decay.
+        warmup_factor (float): The warmup factor.
+        warmup_iter (int): The number of warmup steps.
+        warmup_method (str, optional): The method of warmup, you can choose "linear" or "constant".
+            In linear mode, the multiplication factor starts with warmup_factor in the first epoch
+            and then inreases linearly to reach 1. Defaults to "linear".
+    """
+    exponential_lr = flow.optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)
+    if warmup_iter == 0:
+        logger.warning("warmup iters equals to zero, return ExponentialLR")
+        return exponential_lr
+    warmup_exponential_lr = flow.optim.lr_scheduler.WarmUpLR(
+        exponential_lr,
+        warmup_factor=warmup_factor,
+        warmup_iters=warmup_iter,
+        warmup_method=warmup_method,
+    )
+    return warmup_exponential_lr
+
+
+def WarmupPolynomialLR(
+    optimizer: flow.optim.Optimizer,
+    max_iter: int,
+    warmup_factor: float,
+    warmup_iter: int,
+    end_learning_rate: float = 0.0001,
+    power: float = 1.0,
+    cycle: bool = False,
+    warmup_method: str = "linear",
+):
+    """Create a schedule with a learning rate that decreases as a polynomial decay from
+    the initial lr set in the optimizer to end lr defined by `lr_end`,
+    after a warmup period during which it increases linearly from 0 to the
+    initial lr set in the optimizer.
+
+
+    Args:
+        optimizer (flow.optim.Optimizer): Wrapped optimizer.
+        max_iter (int): Total training iters.
+        warmup_factor (float): The warmup factor.
+        warmup_iter (int): The number of warmup steps.
+        end_learning_rate (float, optional): The final learning rate. Defaults to 0.0001.
+        power (float, optional): The power of polynomial. Defaults to 1.0.
+        cycle (bool, optional): If cycle is True, the scheduler will decay the learning rate
+            every decay steps. Defaults to False.
+        warmup_method (str, optional): The method of warmup, you can choose "linear" or "constant".
+            In linear mode, the multiplication factor starts with warmup_factor in the first
+            epoch and then inreases linearly to reach 1. Defaults to "linear".
+    """
+    polynomial_lr = flow.optim.lr_scheduler.PolynomialLR(
+        optimizer,
+        decay_batch=max_iter,
+        end_learning_rate=end_learning_rate,
+        power=power,
+        cycle=cycle,
+    )
+    if warmup_iter == 0:
+        logger.warning("warmup iters equals to zero, return PolynomialLR")
+        return polynomial_lr
+    warmup_polynomial_lr = flow.optim.lr_scheduler.WarmUpLR(
+        polynomial_lr,
+        warmup_factor=warmup_factor,
+        warmup_iters=warmup_iter,
+        warmup_method=warmup_method,
+    )
+    return warmup_polynomial_lr

libai/tokenizer/__init__.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .build import build_tokenizer
+from .tokenization_bert import BertTokenizer
+from .tokenization_roberta import RobertaTokenizer
+from .tokenization_gpt2 import GPT2Tokenizer
+from .tokenization_t5 import T5Tokenizer
+from .tokenization_base import PreTrainedTokenizer

libai/tokenizer/build.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from libai.config import instantiate
+
+logger = logging.getLogger(__name__)
+
+
+def build_tokenizer(cfg):
+    """Initialize tokenizer."""
+    tokenizer = instantiate(cfg.tokenizer)
+
+    if cfg.append_eod and tokenizer.eod_token is None:
+        if tokenizer.eos_token is not None:
+            tokenizer.eod_token = tokenizer.eos_token
+        else:
+            tokenizer.eod_token = tokenizer.pad_token
+
+    return tokenizer

libai/tokenizer/tokenization_base.py

+# coding=utf-8
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""copy from HuggingFace transformer repo, to tokenize the sentence.
+This class only focus on tokenization, converting token to id and their inverse operation.
+It does not construct inputs using special symbols."""
+
+import copy
+import itertools
+import json
+import logging
+import os
+import unicodedata
+from io import open
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import oneflow as flow
+
+from libai.utils import distributed as dist
+from libai.utils.file_io import PathManager
+from libai.utils.file_utils import cached_path
+
+logger = logging.getLogger(__name__)
+
+
+def _is_whitespace(char):
+    """Checks whether `char` is a whitespace character."""
+    # \t, \n, and \r are technically control characters but we treat them
+    # as whitespace since they are generally considered as such.
+    if char == " " or char == "\t" or char == "\n" or char == "\r":
+        return True
+    cat = unicodedata.category(char)
+    if cat == "Zs":
+        return True
+    return False
+
+
+def _is_control(char):
+    """Checks whether `char` is a control character."""
+    # These are technically control characters but we count them as whitespace
+    # characters.
+    if char == "\t" or char == "\n" or char == "\r":
+        return False
+    cat = unicodedata.category(char)
+    if cat.startswith("C"):
+        return True
+    return False
+
+
+def _is_punctuation(char):
+    """Checks whether `char` is a punctuation character."""
+    cp = ord(char)
+    # We treat all non-letter/number ASCII as punctuation.
+    # Characters such as "^", "$", and "`" are not in the Unicode
+    # Punctuation class but we treat them as punctuation anyways, for
+    # consistency.
+    if (
+        (cp >= 33 and cp <= 47)
+        or (cp >= 58 and cp <= 64)
+        or (cp >= 91 and cp <= 96)
+        or (cp >= 123 and cp <= 126)
+    ):
+        return True
+    cat = unicodedata.category(char)
+    if cat.startswith("P"):
+        return True
+    return False
+
+
+SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json"
+ADDED_TOKENS_FILE = "added_tokens.json"
+TOKENIZER_CONFIG_FILE = "tokenizer_config.json"
+
+
+class PreTrainedTokenizer(object):
+    """
+    Base class for all tokenizers.
+
+    Handle all the shared methods for tokenization and special tokens, methods
+    dowloading/caching/loading pretrained tokenizers as well as adding tokens to the vocabulary.
+    This class also contains the added tokens in a unified way on top of all tokenizers, so we don't
+    have to handle the specific vocabulary augmentation methods of the various underlying
+    dictionary structures (BPE, sentencepiece...).
+
+    Class attributes (overridden by derived classes):
+
+        ``vocab_files_names``: a python ``dict`` with, as keys, the ``__init__`` keyword name of
+        each vocabulary file required by the model, and as associated values, the filename for
+        saving the associated file (string).
+
+        ``pretrained_vocab_files_map``: a python ``dict of dict`` the high-level keys being the
+        ``__init__`` keyword name of each vocabulary file required by the model, the low-level
+        being the `short-cut-names` (string) of the pretrained models with, as associated values,
+        the `url` (string) to the associated pretrained vocabulary file.
+
+        ``max_model_input_sizes``: a python ``dict`` with, as keys, the `short-cut-names` (string)
+        of the pretrained models, and as associated values, the maximum length of the sequence
+        inputs of this model, or None if the model has no maximum input size.
+
+        ``pretrained_init_configuration``: a python ``dict`` with, as keys, the `short-cut-names`
+        (string) of the pretrained models, and as associated values, a dictionnary of specific
+        arguments to pass to the ``__init__`` method of the tokenizer class for this pretrained
+        model when loading the tokenizer with the ``from_pretrained()`` method.
+
+    Args:
+        bos_token (:obj:`str`, `optional`): A special token representing the beginning of a
+            sentence.
+        eos_token (:obj:`str`, `optional`): A special token representing the end of a sentence.
+        unk_token (:obj:`str`, `optional`): A special token representing an out-of-vocabulary token.
+        sep_token (:obj:`str`, `optional`): A special token separating two different sentences in
+            the same input (used by BERT for instance).
+        pad_token (:obj:`str`, `optional`): A special token used to make arrays of tokens the same
+            size for batching purpose.
+            Will then be ignored by attention mechanisms or loss computation.
+        cls_token (:obj:`str`, `optional`): A special token representing the class of the input
+            (used by BERT for instance).
+        mask_token (:obj:`str`, `optional`): A special token representing a masked token (used by
+            masked-language modeling pretraining objectives, like BERT).
+        eod_token (:obj:`str`, `optional`): A special token representing the end of a document.
+            additional_special_tokens (tuple or list of :obj:`str`, `optional`):
+            A tuple or a list of additional special tokens.
+    """
+
+    vocab_files_names = {}
+    pretrained_vocab_files_map = {}
+    pretrained_init_configuration = {}
+    max_model_input_sizes = {}
+
+    SPECIAL_TOKENS_ATTRIBUTES = [
+        "bos_token",
+        "eos_token",
+        "unk_token",
+        "sep_token",
+        "pad_token",
+        "cls_token",
+        "mask_token",
+        "eod_token",
+        "additional_special_tokens",
+    ]
+
+    def __init__(self, verbose=True, **kwargs):
+        self._bos_token = None
+        self._eos_token = None
+        self._unk_token = None
+        self._sep_token = None
+        self._pad_token = None
+        self._cls_token = None
+        self._mask_token = None
+        self._eod_token = None
+        self._additional_special_tokens = []
+        self.verbose = verbose
+
+        # Added tokens - We store this for both slow and fast tokenizers
+        # until the serialization of Fast tokenizers is updated
+        self.added_tokens_encoder: Dict[str, int] = {}
+        self.added_tokens_decoder: Dict[int, str] = {}
+        self.unique_no_split_tokens: List[str] = []
+
+        # inputs and kwargs for saving and re-loading
+        # (see ``from_pretrained`` and ``save_pretrained``)
+        self.init_inputs = ()
+        self.init_kwargs = {}
+
+        # We directly set the hidden value to allow initialization with special tokens
+        # which are not yet in the vocabulary. Necessary for serialization/de-serialization
+        for key, value in kwargs.items():
+            if value is None:
+                continue
+            if key in self.SPECIAL_TOKENS_ATTRIBUTES:
+                if key == "additional_special_tokens":
+                    assert all(
+                        isinstance(t, str) for t in value
+                    ), "One of the tokens is not a string"
+                    setattr(self, key, list(value))
+                elif isinstance(value, str):
+                    setattr(self, key, value)
+                else:
+                    raise TypeError(f"special token {key} has to be str but got: {type(value)}")
+
+    @classmethod
+    def from_pretrained(cls, *inputs, **kwargs):
+        r"""
+        Instantiate a :class:`~PreTrainedTokenizer` (or a derived class) from a
+        predefined tokenizer.
+
+        Args:
+            pretrained_model_name_or_path(`str` or `os.PathLike`):
+                Can be either:
+
+                - a string with the `shortcut name` of a predefined tokenizer to load from cache
+                  or download, e.g.: ``bert-base-uncased``.
+
+                - a path to a `directory` containing vocabulary files required by the tokenizer,
+                  for instance saved using the :func:`~PreTrainedTokenizer.save_pretrained`
+                  method, e.g., ``./my_model_directory/``.
+
+                - (not applicable to all derived classes) a path or url to a single saved
+                  vocabulary file if and only if the tokenizer only requires a single vocabulary
+                  file (e.g. Bert, XLNet), e.g., ``./my_model_directory/vocab.txt``.
+
+            cache_dir: (`optional`) string:
+                Path to a directory in which a downloaded predefined tokenizer vocabulary files
+                should be cached if the standard cache should not be used.
+            force_download: (`optional`) boolean, default False:
+                Force to (re-)download the vocabulary files and override the cached versions if
+                they exist.
+            proxies: (`optional`) dict, default None:
+                A dictionary of proxy servers to use by protocol or endpoint,
+                e.g., {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}.
+                The proxies are used on each request.
+            inputs: (`optional`) positional arguments: will be passed to the
+                Tokenizer ``__init__`` method.
+            kwargs: (`optional`) keyword arguments: will be passed to the
+                Tokenizer ``__init__`` method. Can be used to set special tokens
+                like ``bos_token``, ``eos_token``, ``unk_token``, ``sep_token``,
+                ``pad_token``, ``cls_token``, ``mask_token``, ``additional_special_tokens``.
+                See parameters in the doc string of :class:`~PreTrainedTokenizer`
+                for details.
+
+        Examples:
+
+        .. code-block:: python
+
+            # We can't instantiate directly the base class `PreTrainedTokenizer` so let's
+            # show our examples on a derived class: BertTokenizer
+            # Download vocabulary from S3 and cache.
+            tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+            # If vocabulary files are in a directory (e.g. tokenizer was
+            # saved using `save_pretrained('./test/saved_model/')`)
+            tokenizer = BertTokenizer.from_pretrained('./test/saved_model/')
+            # If the tokenizer uses a single vocabulary file, you can point directly to this file
+            tokenizer = BertTokenizer.from_pretrained('./test/saved_model/my_vocab.txt')
+            # You can link tokens to special vocabulary when instantiating
+            tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', unk_token='<unk>')
+            # You should be sure '<unk>' is in the vocabulary when doing that.
+            # Otherwise use tokenizer.add_special_tokens({'unk_token': '<unk>'}) instead)
+            assert tokenizer.unk_token == '<unk>'
+        """
+        return cls._from_pretrained(*inputs, **kwargs)
+
+    @classmethod
+    def _from_pretrained(cls, pretrained_model_name_or_path, *init_inputs, **kwargs):
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+
+        s3_models = list(cls.max_model_input_sizes.keys())
+        vocab_files = {}
+        init_configuration = {}
+        if pretrained_model_name_or_path in s3_models:
+            # Get the vocabulary from AWS S3 bucket
+            for file_id, map_list in cls.pretrained_vocab_files_map.items():
+                vocab_files[file_id] = map_list[pretrained_model_name_or_path]
+            if (
+                cls.pretrained_init_configuration
+                and pretrained_model_name_or_path in cls.pretrained_init_configuration
+            ):
+                init_configuration = cls.pretrained_init_configuration[
+                    pretrained_model_name_or_path
+                ]
+        else:
+            # Get the vocabulary from local files
+            logger.info(
+                "Model name '{}' not found in model shortcut name list ({}). "
+                "Assuming '{}' is a path or url to a directory containing tokenizer files.".format(
+                    pretrained_model_name_or_path,
+                    ", ".join(s3_models),
+                    pretrained_model_name_or_path,
+                )
+            )
+
+            # Look for the tokenizer main vocabulary files
+            for file_id, file_name in cls.vocab_files_names.items():
+                if os.path.isdir(pretrained_model_name_or_path):
+                    # If a directory is provided we look for the standard filenames
+                    full_file_name = os.path.join(pretrained_model_name_or_path, file_name)
+                else:
+                    # If a path to a file is provided we use it (will only work for non-BPE
+                    # tokenizer using a single vocabulary file)
+                    full_file_name = pretrained_model_name_or_path
+                if not os.path.exists(full_file_name):
+                    logger.info("Didn't find file {}. We won't load it.".format(full_file_name))
+                    full_file_name = None
+                vocab_files[file_id] = full_file_name
+
+            # Look for the additional tokens files
+            additional_files_names = {
+                "added_tokens_file": ADDED_TOKENS_FILE,
+                "special_tokens_map_file": SPECIAL_TOKENS_MAP_FILE,
+                "tokenizer_config_file": TOKENIZER_CONFIG_FILE,
+            }
+
+            # If a path to a file was provided, get the parent directory
+            saved_directory = pretrained_model_name_or_path
+            if os.path.exists(saved_directory) and not os.path.isdir(saved_directory):
+                saved_directory = os.path.dirname(saved_directory)
+
+            for file_id, file_name in additional_files_names.items():
+                full_file_name = os.path.join(saved_directory, file_name)
+                if not os.path.exists(full_file_name):
+                    logger.info("Didn't find file {}. We won't load it.".format(full_file_name))
+                    full_file_name = None
+                vocab_files[file_id] = full_file_name
+
+            if all(full_file_name is None for full_file_name in vocab_files.values()):
+                logger.error(
+                    "Model name '{}' was not found in model name list ({}). "
+                    "We assumed '{}' was a path or url but couldn't find tokenizer files"
+                    "at this path or url.".format(
+                        pretrained_model_name_or_path,
+                        ", ".join(s3_models),
+                        pretrained_model_name_or_path,
+                    )
+                )
+                return None
+
+        # Get files from url, cache, or disk depending on the case
+        try:
+            resolved_vocab_files = {}
+            for file_id, file_path in vocab_files.items():
+                if file_path is None:
+                    resolved_vocab_files[file_id] = None
+                else:
+                    resolved_vocab_files[file_id] = cached_path(
+                        file_path,
+                        cache_dir=cache_dir,
+                        force_download=force_download,
+                        proxies=proxies,
+                    )
+        except EnvironmentError as e:
+            if pretrained_model_name_or_path in s3_models:
+                logger.error("Couldn't reach server to download vocabulary.")
+            else:
+                logger.error(
+                    "Model name '{}' was not found in model name list ({}). "
+                    "We assumed '{}' was a path or url but couldn't find files {} "
+                    "at this path or url.".format(
+                        pretrained_model_name_or_path,
+                        ", ".join(s3_models),
+                        pretrained_model_name_or_path,
+                        str(vocab_files.keys()),
+                    )
+                )
+            raise e
+
+        for file_id, file_path in vocab_files.items():
+            if file_path == resolved_vocab_files[file_id]:
+                logger.info("loading file {}".format(file_path))
+            else:
+                logger.info(
+                    "loading file {} from cache at {}".format(
+                        file_path, resolved_vocab_files[file_id]
+                    )
+                )
+
+        # Prepare tokenizer initialization kwargs
+        # Did we saved some inputs and kwargs to reload ?
+        tokenizer_config_file = resolved_vocab_files.pop("tokenizer_config_file", None)
+        if tokenizer_config_file is not None:
+            init_kwargs = json.load(open(tokenizer_config_file, encoding="utf-8"))
+            saved_init_inputs = init_kwargs.pop("init_inputs", ())
+            if not init_inputs:
+                init_inputs = saved_init_inputs
+        else:
+            init_kwargs = init_configuration
+
+        # Update with newly provided kwargs
+        init_kwargs.update(kwargs)
+
+        # Merge resolved_vocab_files arguments in init_kwargs.
+        added_tokens_file = resolved_vocab_files.pop("added_tokens_file", None)
+        special_tokens_map_file = resolved_vocab_files.pop("special_tokens_map_file", None)
+        for args_name, file_path in resolved_vocab_files.items():
+            if args_name not in init_kwargs:
+                init_kwargs[args_name] = file_path
+
+        if special_tokens_map_file is not None:
+            special_tokens_map = json.load(open(special_tokens_map_file, encoding="utf-8"))
+            for key, value in special_tokens_map.items():
+                if key not in init_kwargs:
+                    init_kwargs[key] = value
+
+        # Instantiate tokenizer.
+        tokenizer = cls(*init_inputs, **init_kwargs)
+
+        # Save inputs and kwargs for saving and re-loading with ``save_pretrained``
+        tokenizer.init_inputs = init_inputs
+        tokenizer.init_kwargs = init_kwargs
+
+        # Add supplementary tokens.
+        special_tokens = tokenizer.all_special_tokens
+        if added_tokens_file is not None:
+            with open(added_tokens_file, encoding="utf-8") as added_tokens_handle:
+                added_tok_encoder = json.load(added_tokens_handle)
+
+            # Sort added tokens by index
+            added_tok_encoder_sorted = list(sorted(added_tok_encoder.items(), key=lambda x: x[1]))
+
+            for token, index in added_tok_encoder_sorted:
+                assert index == len(tokenizer), (
+                    f"Non-consecutive added token '{token}' found. "
+                    f"Should have index {len(tokenizer)} but has index {index} in saved vocabulary."
+                )
+                tokenizer.add_tokens(token, special_tokens=bool(token in special_tokens))
+
+        # Check all our special tokens are registered as "no split" token
+        # (we don't cut them) and are in the vocab
+        added_tokens = tokenizer.sanitize_special_tokens()
+        if added_tokens:
+            logger.warning(
+                "Special tokens have been added in the vocabulary,"
+                "make sure the associated word embedding are fine-tuned or trained."
+            )
+
+        return tokenizer
+
+    def save_pretrained(self, save_directory):
+        """
+        Save the tokenizer vocabulary files together with:
+
+            - added tokens,
+            - special-tokens-to-class-attributes-mapping,
+            - tokenizer instantiation positional and keywords inputs (e.g. do_lower_case for Bert).
+
+        This won't save modifications other than ``added tokens`` and ``special token mapping``,
+        you may have applied to the tokenizer after the instantiation (e.g. modifying
+        tokenizer.do_lower_case after creation).
+        This method make sure the full tokenizer can then be re-loaded using the
+        :func:`~PreTrainedTokenizer.from_pretrained` class method.
+        """
+        if not PathManager.isdir(save_directory):
+            logger.error("Saving directory ({}) should be a directory".format(save_directory))
+            return
+        PathManager.mkdirs(save_directory)
+
+        special_tokens_map_file = os.path.join(save_directory, SPECIAL_TOKENS_MAP_FILE)
+        added_tokens_file = os.path.join(save_directory, ADDED_TOKENS_FILE)
+        tokenizer_config_file = os.path.join(save_directory, TOKENIZER_CONFIG_FILE)
+
+        tokenizer_config = copy.deepcopy(self.init_kwargs)
+        if len(self.init_inputs) > 0:
+            tokenizer_config["init_inputs"] = copy.deepcopy(self.init_inputs)
+        for file_id in self.vocab_files_names.keys():
+            tokenizer_config.pop(file_id, None)
+
+        with open(tokenizer_config_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(tokenizer_config, ensure_ascii=False))
+
+        with open(special_tokens_map_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self.special_tokens_map, ensure_ascii=False))
+
+        added_vocab = self.get_added_vocab()
+        if added_vocab:
+            with open(added_tokens_file, "w", encoding="utf-8") as f:
+                out_str = json.dumps(added_vocab, ensure_ascii=False)
+                f.write(out_str)
+
+        vocab_files = self.save_vocabulary(save_directory)
+
+        return vocab_files + (special_tokens_map_file, added_tokens_file)
+
+    def save_vocabulary(self, save_directory):
+        """Save the tokenizer vocabulary to a directory. This method does *NOT* save added tokens
+        and special token mappings.
+        Please use :func:`~PreTrainedTokenizer.save_pretrained` to save the
+        full Tokenizer state if you want to reload it using the
+        :func:`~PreTrainedTokenizer.from_pretrained` class method.
+        """
+        raise NotImplementedError
+
+    @property
+    def vocab_size(self) -> int:
+        """Size of the base vocabulary (without the added tokens)."""
+        raise NotImplementedError
+
+    def padded_vocab_size(self, multiple=1) -> int:
+        """Padded the vocabulary with dummy tokens and return the new size."""
+        vocab_size = len(self)
+        while vocab_size % multiple != 0:
+            vocab_size += 1
+        return vocab_size
+
+    def __len__(self):
+        """Size of the full vocabulary with the added tokens."""
+        return self.vocab_size + len(self.added_tokens_encoder)
+
+    def get_vocab(self) -> Dict[str, int]:
+        """
+        Returns the vocabulary as a dictionary of token to index.
+        :obj:`tokenizer.get_vocab()[token]` is equivalent to
+        :obj:`tokenizer.convert_tokens_to_ids(token)`
+        when :obj:`token` is in the vocab.
+
+        Returns:
+            :obj:`Dict[str, int]`: The vocabulary.
+        """
+        raise NotImplementedError
+
+    def get_added_vocab(self) -> Dict[str, int]:
+        """
+        Returns the added tokens in the vocabulary as a dictionary of token to index.
+
+        Returns:
+            :obj:`Dict[str, int]`: The added tokens.
+        """
+        return self.added_tokens_encoder
+
+    def add_tokens(self, new_tokens: Union[str, List[str]], special_tokens: bool = False) -> int:
+        """
+        Add a list of new tokens to the tokenizer class. If the new tokens are not in the
+        vocabulary, they are added to it with indices starting from the length of
+        the current vocabulary.
+
+        .. Note::
+            When adding new tokens to the vocabulary, you should make sure to also resize
+            the token embedding matrix of the model so that its embedding matrix matches
+            the tokenizer.
+            In order to do that, please use the
+            :meth:`~PreTrainedModel.resize_token_embeddings` method.
+
+        Args:
+            new_tokens (:obj:`str`, or a list of `str`):
+                Tokens are only added if they are not already in the vocabulary.
+            special_tokens (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Can be used to specify if the token is a special token. This mostly change
+                the normalization behavior
+                (special tokens like CLS or [MASK] are usually not lower-cased for instance).
+
+        Returns:
+            :obj:`int`: Number of tokens added to the vocabulary.
+
+        Examples:
+
+        .. code-block:: python
+
+            # Let's see how to increase the vocabulary of Bert model and tokenizer
+            tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
+            model = BertModel.from_pretrained('bert-base-uncased')
+            num_added_toks = tokenizer.add_tokens(['new_tok1', 'my_new-tok2'])
+            print('We have added', num_added_toks, 'tokens')
+             # Notice: resize_token_embeddings expect to receive the full size of the new
+             # vocabulary, i.e., the length of the tokenizer.
+            model.resize_token_embeddings(len(tokenizer))
+        """
+        if not new_tokens:
+            return 0
+
+        if not isinstance(new_tokens, (list, tuple)):
+            new_tokens = [new_tokens]
+
+        tokens_to_add = []
+        for token in new_tokens:
+            if not isinstance(token, str):
+                raise TypeError(f"Token {token} is not a string but a {type(token)}.")
+            if not special_tokens and hasattr(self, "do_lower_case") and self.do_lower_case:
+                token = token.lower()
+            if (
+                token != self.unk_token
+                and self.convert_tokens_to_ids(token) == self.convert_tokens_to_ids(self.unk_token)
+                and token not in tokens_to_add
+            ):
+                tokens_to_add.append(token)
+                if self.verbose:
+                    logger.info(f"Adding {token} to the vocabulary")
+
+        added_tok_encoder = dict((tok, len(self) + i) for i, tok in enumerate(tokens_to_add))
+        added_tok_decoder = {v: k for k, v in added_tok_encoder.items()}
+        self.added_tokens_encoder.update(added_tok_encoder)
+        self.added_tokens_decoder.update(added_tok_decoder)
+
+        if special_tokens:
+            self.unique_no_split_tokens = sorted(
+                set(self.unique_no_split_tokens).union(set(new_tokens))
+            )
+        else:
+            self.unique_no_split_tokens = sorted(
+                set(self.unique_no_split_tokens).union(set(tokens_to_add))
+            )
+
+        return len(tokens_to_add)
+
+    def sanitize_special_tokens(self) -> int:
+        """
+        Make sure that all the special tokens attributes of the tokenizer
+        (:obj:`tokenizer.mask_token`, :obj:`tokenizer.cls_token`, etc.)
+        are in the vocabulary.
+
+        Add the missing ones to the vocabulary if needed.
+
+        Return:
+            :obj:`int`: The number of tokens added in the vocaulary during the operation.
+        """
+        return self.add_tokens(self.all_special_tokens, special_tokens=True)
+
+    def add_special_tokens(self, special_tokens_dict: Dict[str, str]) -> int:
+        """
+        Add a dictionary of special tokens (eos, pad, cls, etc.) to the encoder and link them to
+        class attributes. If special tokens are NOT in the vocabulary, they are added to it
+        (indexed starting from the last index of the current vocabulary).
+
+        .. Note::
+            When adding new tokens to the vocabulary, you should make sure to also resize the
+            token embedding matrix of the model so that its embedding matrix matches the tokenizer.
+            In order to do that, please use the
+            :meth:`~PreTrainedModel.resize_token_embeddings` method.
+
+        Using :obj:`add_special_tokens` will ensure your special tokens can be used in several ways:
+        - Special tokens are carefully handled by the tokenizer (they are never split).
+        - You can easily refer to special tokens using tokenizer class attributes like
+        :obj:`tokenizer.cls_token`. This makes it easy to develop model-agnostic training and
+        fine-tuning scripts.
+        When possible, special tokens are already registered for provided pretrained models
+        (for instance :class:`~BertTokenizer` :obj:`cls_token` is already registered
+        to be :obj`'[CLS]'` and XLM's one is also registered to be :obj:`'</s>'`).
+
+        Args:
+            special_tokens_dict (dictionary `str` to `str`):
+                Keys should be in the list of predefined special attributes: [``bos_token``,
+                ``eos_token``, ``unk_token``, ``sep_token``, ``pad_token``,
+                ``cls_token``, ``mask_token``,
+                ``additional_special_tokens``].
+                Tokens are only added if they are not already in the vocabulary (tested by
+                checking if the tokenizer assign the index of the ``unk_token`` to them).
+
+        Returns:
+            :obj:`int`: Number of tokens added to the vocabulary.
+
+        Examples:
+
+        .. code-block:: python
+
+            # Let's see how to add a new classification token to GPT-2
+            tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+            model = GPT2Model.from_pretrained('gpt2')
+            special_tokens_dict = {'cls_token': '<CLS>'}
+            num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)
+            print('We have added', num_added_toks, 'tokens')
+            # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary,
+            # i.e., the length of the tokenizer.
+            model.resize_token_embeddings(len(tokenizer))
+            assert tokenizer.cls_token == '<CLS>'
+        """
+        if not special_tokens_dict:
+            return 0
+
+        added_tokens = 0
+        for key, value in special_tokens_dict.items():
+            assert key in self.SPECIAL_TOKENS_ATTRIBUTES, f"Key {key} is not a special token"
+
+            if self.verbose:
+                logger.info(f"Assigning {value} to the {key} key of the tokenizer")
+            setattr(self, key, value)
+
+            if key == "additional_special_tokens":
+                assert isinstance(value, (list, tuple)) and all(
+                    isinstance(t, str) for t in value
+                ), f"Tokens {value} for key {key} should all be a string"
+                added_tokens += self.add_tokens(value, special_tokens=True)
+            else:
+                assert isinstance(value, str), f"Token {value} for key {key} should be a string"
+                added_tokens += self.add_tokens([value], special_tokens=True)
+
+        return added_tokens
+
+    def tokenize(self, text: str, **kwargs) -> List[str]:
+        """
+        Converts a string in a sequence of tokens, using the tokenizer.
+        Split in words for word-based vocabulary or sub-words for sub-word-based vocabularies
+        (BPE/SentencePieces/WordPieces). Take care of added tokens.
+
+        Args:
+            text (:obj:`str`):
+                The sequence to be encoded.
+            **kwargs (additional keyword arguments):
+                Passed along to the model-specific ``prepare_for_tokenization``
+                preprocessing method.
+
+        Returns:
+            :obj:`List[str]`: The list of tokens.
+        """
+
+        def split_on_token(tok, text):
+            result = []
+            split_text = text.split(tok)
+            for i, sub_text in enumerate(split_text):
+                sub_text = sub_text.strip()
+                if i == 0 and not sub_text:
+                    result += [tok]
+                elif i == len(split_text) - 1:
+                    if sub_text:
+                        result += [sub_text]
+                    else:
+                        pass
+                else:
+                    if sub_text:
+                        result += [sub_text]
+                    result += [tok]
+            return result
+
+        def split_on_tokens(tok_list, text):
+            if not text:
+                return []
+            if not tok_list:
+                return self._tokenize(text, **kwargs)
+
+            tokenized_text = []
+            text_list = [text]
+            for tok in tok_list:
+                tokenized_text = []
+                for sub_text in text_list:
+                    if sub_text not in self.unique_no_split_tokens:
+                        tokenized_text += split_on_token(tok, sub_text)
+                    else:
+                        tokenized_text += [sub_text]
+                text_list = tokenized_text
+            return list(
+                itertools.chain.from_iterable(
+                    (
+                        self._tokenize(token)
+                        if token not in self.unique_no_split_tokens
+                        else [token]
+                        for token in tokenized_text
+                    )
+                )
+            )
+
+        no_split_token = self.unique_no_split_tokens
+        tokenized_text = split_on_tokens(no_split_token, text)
+        return tokenized_text
+
+    def _tokenize(self, text, **kwargs):
+        """
+        Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for
+        word-based vocabulary or sub-words for sub-word-based vocabularies
+        (BPE/SentencePieces/WordPieces).
+        Do NOT take care of added tokens.
+        """
+        raise NotImplementedError
+
+    def convert_tokens_to_ids(self, tokens: Union[str, List[str]]) -> Union[int, List[int]]:
+        """Converts a token string (or a sequence of tokens) in a single integer id
+        (or a sequence of ids), using the vocabulary.
+        """
+        if tokens is None:
+            return None
+
+        if isinstance(tokens, str):
+            return self._convert_token_to_id_with_added_voc(tokens)
+
+        if len(tokens) > 0 and isinstance(tokens[0], list):
+            ids = []
+            for ts in tokens:
+                ids_x = []
+                for token in ts:
+                    ids_x.append(self._convert_token_to_id_with_added_voc(token))
+                ids.append(ids_x)
+            return ids
+
+        ids = []
+        for token in tokens:
+            ids.append(self._convert_token_to_id_with_added_voc(token))
+        return ids
+
+    def convert_to_tensors(self, token_ids, return_tensors=None, is_global=False, **kwargs):
+        if return_tensors is None:
+            return_token_ids = token_ids
+        elif return_tensors == "of":
+            if not is_global:
+                return_token_ids = flow.tensor(token_ids, dtype=flow.long)
+            elif is_global:
+                sbp = kwargs.get("sbp", dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]))
+                placement = kwargs.get(
+                    "placement", flow.placement("cuda", list(range(dist.get_world_size())))
+                )
+                return_token_ids = flow.tensor(
+                    token_ids, sbp=sbp, placement=placement, dtype=flow.long
+                )
+        elif return_tensors == "np":
+            return_token_ids = np.array(token_ids, dtype=np.int64)
+        return return_token_ids
+
+    def _convert_token_to_id_with_added_voc(self, token):
+        if token is None:
+            return None
+
+        if token in self.added_tokens_encoder:
+            return self.added_tokens_encoder[token]
+        return self._convert_token_to_id(token)
+
+    def _convert_token_to_id(self, token):
+        raise NotImplementedError
+
+    def encode(self, text, return_tensors=None, is_global=False, **kwargs):
+        if isinstance(text, str):
+            tokens = self.tokenize(text)
+            token_ids = self.convert_tokens_to_ids(tokens)
+            token_ids = self.build_inputs_with_special_tokens(token_ids)
+            token_ids = self.convert_to_tensors(
+                token_ids, return_tensors=return_tensors, is_global=is_global, **kwargs
+            )
+            return token_ids
+        elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], str):
+            tokens = [self.tokenize(t) for t in text]
+            token_ids_list = self.convert_tokens_to_ids(tokens)
+            token_ids_list = [
+                self.build_inputs_with_special_tokens(token_ids) for token_ids in token_ids_list
+            ]
+            token_ids_list = self.convert_to_tensors(
+                token_ids_list, return_tensors=return_tensors, is_global=is_global, **kwargs
+            )
+            return token_ids_list
+        elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int):
+            return text
+        else:
+            raise ValueError(
+                "Input is not valid. Should be a string, a list/tuple of strings or "
+                "a list/tuple of integers."
+            )
+
+    def convert_ids_to_tokens(
+        self, ids: Union[int, List[int]], skip_special_tokens: bool = False
+    ) -> Union[str, List[str]]:
+        """
+        Converts a single index or a sequence of indices in a token or a sequence of tokens,
+        using the vocabulary and added tokens.
+
+        Args:
+            ids (:obj:`int` or :obj:`List[int]`):
+                The token id (or token ids) to convert to tokens.
+            skip_special_tokens (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to remove special tokens in the decoding.
+        Returns:
+            :obj:`str` or :obj:`List[str]`: The decoded token(s).
+        """
+        if isinstance(ids, int):
+            if ids in self.added_tokens_decoder:
+                return self.added_tokens_decoder[ids]
+            else:
+                return self._convert_id_to_token(ids)
+        tokens = []
+        for index in ids:
+            if skip_special_tokens and index in self.all_special_ids:
+                continue
+            if index in self.added_tokens_decoder:
+                tokens.append(self.added_tokens_decoder[index])
+            else:
+                tokens.append(self._convert_id_to_token(index))
+        return tokens
+
+    def _convert_id_to_token(self, index: int) -> str:
+        raise NotImplementedError
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """
+        Converts a sequence of tokens to a single string. The most simple way to do it is
+        ``" ".join(tokens)`` but we often want to remove sub-word tokenization artifacts
+        at the same time.
+
+        Args:
+            tokens (:obj:`List[str]`): The token to join in a string.
+
+        Returns:
+            :obj:`str`: The joined tokens.
+        """
+        return " ".join(tokens)
+
+    def decode(
+        self,
+        token_ids,
+        skip_special_tokens=False,
+        clean_up_tokenization_spaces=True,
+        spaces_between_special_tokens: bool = True,
+    ):
+        """
+        Converts a sequence of ids (integer) in a string, using the tokenizer and vocabulary
+        with options to remove special tokens and clean up tokenization spaces.
+        Similar to doing ``self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))``.
+
+        Args:
+            token_ids: list of tokenized input ids. Can be obtained using the `encode` or
+            `encode_plus` methods.
+            skip_special_tokens: if set to True, will replace special tokens.
+            clean_up_tokenization_spaces: if set to True, will clean up the tokenization spaces.
+        """
+        # Convert inputs to python lists
+        if isinstance(token_ids, flow.Tensor):
+            token_ids = token_ids.tolist()
+
+        filtered_tokens = self.convert_ids_to_tokens(
+            token_ids, skip_special_tokens=skip_special_tokens
+        )
+
+        # To avoid mixing byte-level and unicode for byte-level BPT
+        # we need to build string separately for added tokens and byte-level tokens
+        # cf. https://github.com/huggingface/transformers/issues/1133
+        sub_texts = []
+        current_sub_text = []
+        for token in filtered_tokens:
+            if skip_special_tokens and token in self.all_special_ids:
+                continue
+            if token in self.added_tokens_encoder:
+                if current_sub_text:
+                    sub_texts.append(self.convert_tokens_to_string(current_sub_text))
+                    current_sub_text = []
+                sub_texts.append(token)
+            else:
+                current_sub_text.append(token)
+        if current_sub_text:
+            sub_texts.append(self.convert_tokens_to_string(current_sub_text))
+
+        if spaces_between_special_tokens:
+            text = " ".join(sub_texts)
+        else:
+            text = "".join(sub_texts)
+
+        if clean_up_tokenization_spaces:
+            clean_text = self.clean_up_tokenization(text)
+            return clean_text
+        else:
+            return text
+
+    @property
+    def bos_token(self) -> str:
+        """
+        :obj:`str`: Beginning of sentence token. Log an error if used while not having been set.
+        """
+        if self._bos_token is None and self.verbose:
+            logger.error("Using bos_token, but it is not set yet.")
+            return None
+        return str(self._bos_token)
+
+    @property
+    def eos_token(self) -> str:
+        """
+        :obj:`str`: End of sentence token. Log an error if used while not having been set.
+        """
+        if self._eos_token is None and self.verbose:
+            logger.error("Using eos_token, but it is not set yet.")
+            return None
+        return str(self._eos_token)
+
+    @property
+    def unk_token(self) -> str:
+        """
+        :obj:`str`: Unknown token. Log an error if used while not having been set.
+        """
+        if self._unk_token is None and self.verbose:
+            logger.error("Using unk_token, but it is not set yet.")
+            return None
+        return str(self._unk_token)
+
+    @property
+    def sep_token(self) -> str:
+        """
+        :obj:`str`: Separation token, to separate context and query in an input sequence.
+        Log an error if used while not having been set.
+        """
+        if self._sep_token is None and self.verbose:
+            logger.error("Using sep_token, but it is not set yet.")
+            return None
+        return str(self._sep_token)
+
+    @property
+    def pad_token(self) -> str:
+        """
+        :obj:`str`: Padding token. Log an error if used while not having been set.
+        """
+        if self._pad_token is None and self.verbose:
+            logger.error("Using pad_token, but it is not set yet.")
+            return None
+        return str(self._pad_token)
+
+    @property
+    def cls_token(self) -> str:
+        """
+        :obj:`str`: Classification token, to extract a summary of an input sequence leveraging
+        self-attention along the full depth of the model.
+        Log an error if used while not having been set.
+        """
+        if self._cls_token is None and self.verbose:
+            logger.error("Using cls_token, but it is not set yet.")
+            return None
+        return str(self._cls_token)
+
+    @property
+    def mask_token(self) -> str:
+        """
+        :obj:`str`: Mask token, to use when training a model with masked-language modeling.
+        Log an error if used while not having been set.
+        """
+        if self._mask_token is None and self.verbose:
+            logger.error("Using mask_token, but it is not set yet.")
+            return None
+        return str(self._mask_token)
+
+    @property
+    def eod_token(self) -> str:
+        """
+        :obj:`str`: End of document token. Log an error if used while not having been set.
+        """
+        if self._eod_token is None and self.verbose:
+            logger.error("Using eod_token, but it is not set yet.")
+            return None
+        return str(self._eod_token)
+
+    @property
+    def start_token(self) -> str:
+        """
+        :obj:`str`: Start token of sentence. Common name for bos_token and cls_token.
+        """
+        if self._bos_token is not None and self._cls_token is not None:
+            if self._bos_token == self._cls_token:
+                return str(self._bos_token)
+            else:
+                logger.error("Conflict between bos_token and cls_token.")
+                return None
+        elif self._bos_token is None and self._cls_token is not None:
+            return str(self._cls_token)
+        elif self._bos_token is not None and self._cls_token is None:
+            return str(self._bos_token)
+        else:
+            logger.error("Using start_token, but it is not set yet.")
+            return None
+
+    @property
+    def end_token(self) -> str:
+        """
+        :obj:`str`: End token of sentence. Common name for eos_token and sep_token.
+        Note: eod_token is not considered, because it is often same with eos_token.
+        """
+        if self._eos_token is not None and self._sep_token is not None:
+            if self._eos_token == self._sep_token:
+                return str(self._eos_token)
+            else:
+                logger.error("Conflict between eos_token and _sep_token.")
+                return None
+        elif self._eos_token is None and self._sep_token is not None:
+            return str(self._sep_token)
+        elif self._eos_token is not None and self._sep_token is None:
+            return str(self._eos_token)
+        else:
+            logger.error("Using end_token, but it is not set yet.")
+            return None
+
+    @property
+    def additional_special_tokens(self) -> List[str]:
+        """
+        :obj:`List[str]`: All the additional special tokens you may want to use.
+        Log an error if used while not having been set.
+        """
+        if self._additional_special_tokens is None and self.verbose:
+            logger.error("Using additional_special_tokens, but it is not set yet.")
+            return None
+        return [str(tok) for tok in self._additional_special_tokens]
+
+    @bos_token.setter
+    def bos_token(self, value):
+        self._bos_token = value
+
+    @eos_token.setter
+    def eos_token(self, value):
+        self._eos_token = value
+
+    @unk_token.setter
+    def unk_token(self, value):
+        self._unk_token = value
+
+    @sep_token.setter
+    def sep_token(self, value):
+        self._sep_token = value
+
+    @pad_token.setter
+    def pad_token(self, value):
+        self._pad_token = value
+
+    @cls_token.setter
+    def cls_token(self, value):
+        self._cls_token = value
+
+    @mask_token.setter
+    def mask_token(self, value):
+        self._mask_token = value
+
+    @eod_token.setter
+    def eod_token(self, value):
+        self._eod_token = value
+
+    @additional_special_tokens.setter
+    def additional_special_tokens(self, value):
+        self._additional_special_tokens = value
+
+    @property
+    def bos_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the beginning of sentence token in the vocabulary.
+        Returns :obj:`None` if the token has not been set.
+        """
+        if self._bos_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.bos_token)
+
+    @property
+    def eos_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the end of sentence token in the vocabulary.
+        Returns :obj:`None` if the token has not been set.
+        """
+        if self._eos_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.eos_token)
+
+    @property
+    def unk_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the unknown token in the vocabulary.
+        Returns :obj:`None` if the token has not been set.
+        """
+        if self._unk_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.unk_token)
+
+    @property
+    def sep_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the separation token in the vocabulary,
+        to separate context and query in an input sequence.
+        Returns :obj:`None` if the token has not been set.
+        """
+        if self._sep_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.sep_token)
+
+    @property
+    def pad_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the padding token in the vocabulary.
+        Returns :obj:`None` if the token has not been set.
+        """
+        if self._pad_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.pad_token)
+
+    @property
+    def cls_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the classification token in the vocabulary,
+        to extract a summary of an input sequence leveraging self-attention
+        along the full depth of the model.
+        Returns :obj:`None` if the token has not been set.
+        """
+        if self._cls_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.cls_token)
+
+    @property
+    def mask_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the mask token in the vocabulary, used when training a
+        model with masked-language modeling. Returns :obj:`None` if the token has not been set.
+        """
+        if self._mask_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.mask_token)
+
+    @property
+    def eod_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the end of document token in the vocabulary.
+        Returns :obj:`None` if the token has not been set.
+        """
+        if self._eod_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.eod_token)
+
+    @property
+    def start_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the start token in the vocabulary.
+        Returns :obj:`None` if the token has not been set.
+        """
+        start_token = self.start_token
+        if start_token is None:
+            return None
+        else:
+            return self.convert_tokens_to_ids(start_token)
+
+    @property
+    def end_token_id(self) -> Optional[int]:
+        """
+        :obj:`Optional[int]`: Id of the end token in the vocabulary.
+        Returns :obj:`None` if the token has not been set.
+        """
+        end_token = self.end_token
+        if end_token is None:
+            return None
+        else:
+            return self.convert_tokens_to_ids(end_token)
+
+    @property
+    def additional_special_tokens_ids(self) -> List[int]:
+        """
+        :obj:`List[int]`: Ids of all the additional special tokens in the vocabulary.
+        Log an error if used while not having been set.
+        """
+        return self.convert_tokens_to_ids(self.additional_special_tokens)
+
+    @property
+    def special_tokens_map(self) -> Dict[str, Union[str, List[str]]]:
+        """
+        A dictionary mapping special token class attributes
+        (:obj:`cls_token`, :obj:`unk_token`, etc.) to their values
+        (:obj:`'<unk>'`, :obj:`'<cls>'`, etc.).
+        """
+        set_attr = {}
+        for attr in self.SPECIAL_TOKENS_ATTRIBUTES:
+            attr_value = getattr(self, "_" + attr)
+            if attr_value:
+                set_attr[attr] = attr_value
+        return set_attr
+
+    @property
+    def all_special_tokens(self) -> List[str]:
+        """
+        :obj:`List[str]`: All the special tokens
+        (:obj:`'<unk>'`, :obj:`'<cls>'`, etc.) mapped to class attributes.
+        """
+        all_toks = []
+        set_attr = self.special_tokens_map
+        for attr_value in set_attr.values():
+            all_toks = all_toks + (
+                list(attr_value) if isinstance(attr_value, (list, tuple)) else [attr_value]
+            )
+        all_toks = list(set(all_toks))
+        return all_toks
+
+    @property
+    def all_special_ids(self) -> List[int]:
+        """
+        :obj:`List[int]`: List the ids of the special tokens
+        (:obj:`'<unk>'`, :obj:`'<cls>'`, etc.) mapped to class attributes.
+        """
+        all_toks = self.all_special_tokens
+        all_ids = list(self.convert_tokens_to_ids(all_toks))
+        return all_ids
+
+    @staticmethod
+    def clean_up_tokenization(out_string):
+        """Clean up a list of simple English tokenization artifacts like spaces before
+        punctuations and abbreviated forms.
+        """
+        out_string = (
+            out_string.replace(" .", ".")
+            .replace(" ?", "?")
+            .replace(" !", "!")
+            .replace(" ,", ",")
+            .replace(" ' ", "'")
+            .replace(" n't", "n't")
+            .replace(" 'm", "'m")
+            .replace(" do not", " don't")
+            .replace(" 's", "'s")
+            .replace(" 've", "'ve")
+            .replace(" 're", "'re")
+        )
+        return out_string

libai/tokenizer/tokenization_bert.py

+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tokenization classes for bert (wordpieces)."""
+
+import collections
+import logging
+import os
+import re
+import unicodedata
+from io import open
+from typing import List, Optional
+
+from .tokenization_base import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
+
+logger = logging.getLogger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "bert-base-uncased": "https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt",
+        "bert-large-uncased": "https://huggingface.co/bert-large-uncased/resolve/main/vocab.txt",
+        "bert-base-cased": "https://huggingface.co/bert-base-cased/resolve/main/vocab.txt",
+        "bert-large-cased": "https://huggingface.co/bert-large-cased/resolve/main/vocab.txt",
+        "bert-base-chinese": "https://huggingface.co/bert-base-chinese/resolve/main/vocab.txt",
+    }
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "bert-base-uncased": 512,
+    "bert-large-uncased": 512,
+    "bert-base-cased": 512,
+    "bert-large-cased": 512,
+    "bert-base-chinese": 512,
+}
+
+PRETRAINED_INIT_CONFIGURATION = {
+    "bert-base-uncased": {"do_lower_case": True},
+    "bert-large-uncased": {"do_lower_case": True},
+    "bert-base-cased": {"do_lower_case": False},
+    "bert-large-cased": {"do_lower_case": False},
+    "bert-base-chinese": {"do_lower_case": False},
+}
+
+
+def load_vocab(vocab_file):
+    """Loads a vocabulary file into a dictionary."""
+    vocab = collections.OrderedDict()
+    with open(vocab_file, "r", encoding="utf-8") as reader:
+        tokens = reader.readlines()
+    for index, token in enumerate(tokens):
+        token = token.rstrip("\n")
+        vocab[token] = index
+    return vocab
+
+
+def whitespace_tokenize(text):
+    """Runs basic whitespace cleaning and splitting on a piece of text."""
+    text = text.strip()
+    if not text:
+        return []
+    tokens = text.split()
+    return tokens
+
+
+def _is_chinese_substr(char):
+    return re.findall("##[\u4E00-\u9FA5]", char)
+
+
+class BertTokenizer(PreTrainedTokenizer):
+    """
+    Construct a BERT tokenizer. Based on WordPiece.
+
+    Args:
+        vocab_file (:obj:`str`):
+            Path to a one-wordpiece-per-line vocabulary file.
+        do_lower_case (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether to lower case the input.
+            Only has an effect when do_basic_tokenize=True.
+        do_basic_tokenize (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether to do basic tokenization before wordpiece.
+        never_split (:obj:`Iterable`, `optional`):
+            List of tokens which will never be split during tokenization.
+            Only has an effect when do_basic_tokenize=True.
+        tokenize_chinese_chars (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether to tokenize Chinese characters.
+            This should likely be deactivated for Japanese,
+            see: https://github.com/huggingface/pytorch-pretrained-BERT/issues/328.
+        do_chinese_wwm (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to do whole word masking for Chinese.
+            Chinese sentence will be segmented by a third-party tool first.
+            Each substr will be added '##' prefix and its index will be calucated by
+            id(##A) = id(A) + vocab_size.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    pretrained_init_configuration = PRETRAINED_INIT_CONFIGURATION
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+
+    def __init__(
+        self,
+        vocab_file,
+        do_lower_case=True,
+        do_basic_tokenize=True,
+        never_split=None,
+        unk_token="[UNK]",
+        sep_token="[SEP]",
+        pad_token="[PAD]",
+        cls_token="[CLS]",
+        mask_token="[MASK]",
+        tokenize_chinese_chars=True,
+        do_chinese_wwm=False,
+        add_bos_token=False,
+        **kwargs,
+    ):
+        super(BertTokenizer, self).__init__(
+            unk_token=unk_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            **kwargs,
+        )
+        if not os.path.isfile(vocab_file):
+            raise ValueError(
+                "Can't find a vocabulary file at path '{}'. To load the "
+                "vocabulary from a Google pretrained model use "
+                "`tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
+                    vocab_file
+                )
+            )
+        self.vocab = load_vocab(vocab_file)
+        self.ids_to_tokens = collections.OrderedDict(
+            [(ids, tok) for tok, ids in self.vocab.items()]
+        )
+        self.do_basic_tokenize = do_basic_tokenize
+        if do_basic_tokenize:
+            if do_chinese_wwm:
+                self.basic_tokenizer = BasicTokenizerWithChineseWWM(
+                    do_lower_case=do_lower_case,
+                    never_split=never_split,
+                    tokenize_chinese_chars=tokenize_chinese_chars,
+                )
+            else:
+                self.basic_tokenizer = BasicTokenizer(
+                    do_lower_case=do_lower_case,
+                    never_split=never_split,
+                    tokenize_chinese_chars=tokenize_chinese_chars,
+                )
+        self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=self.unk_token)
+        self.add_bos_token = add_bos_token
+
+    @property
+    def vocab_size(self):
+        return len(self.vocab)
+
+    def get_vocab(self):
+        return dict(self.vocab, **self.added_tokens_encoder)
+
+    def _tokenize(self, text):
+        split_tokens = []
+        if self.do_basic_tokenize:
+            for token in self.basic_tokenizer.tokenize(text, never_split=self.all_special_tokens):
+                for sub_token in self.wordpiece_tokenizer.tokenize(token):
+                    split_tokens.append(sub_token)
+        else:
+            split_tokens = self.wordpiece_tokenizer.tokenize(text)
+        return split_tokens
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab.
+        For Chinese substr, id = vocab_size + id(substr.remove(##)).
+        """
+        index = self.vocab.get(token, self.vocab.get(self.unk_token))
+        return index
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab.
+        For Chinese substr, id = vocab_size + id(substr.remove(##)).
+        """
+        token = self.ids_to_tokens.get(index, self.unk_token)
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) to a single string."""
+        out_string = " ".join(tokens).replace(" ##", "").strip()
+        return out_string
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """Add special tokens to a sequence or a pair of sequence.
+        BERT format sentence input:
+
+        - single sequence: [CLS] tokens_a [SEP]
+        - pair of sequences: [CLS] tokens_a [SEP] tokens_b [SEP]
+
+        Args:
+            token_ids_0 (List[int]): The token ids of sentence 0.
+            token_ids_1 (List[int], optional): The token ids of sentence 1. Defaults to None.
+
+        Returns:
+            :obj:`List[str]`: The sequence after adding special toekens.
+        """
+        if self.add_bos_token:
+            cls = [self.cls_token_id]
+            sep = [self.sep_token_id]
+        else:
+            cls = []
+            sep = []
+
+        if token_ids_1 is None:
+            return cls + token_ids_0 + sep
+
+        return cls + token_ids_0 + sep + token_ids_1 + sep
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        """Save the tokenizer vocabulary to a directory or file."""
+        index = 0
+        if os.path.isdir(save_directory):
+            vocab_file = os.path.join(
+                save_directory,
+                (filename_prefix + "-" if filename_prefix else "")
+                + VOCAB_FILES_NAMES["vocab_file"],
+            )
+        else:
+            vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
+        with open(vocab_file, "w", encoding="utf-8") as writer:
+            for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
+                if index != token_index:
+                    logger.warning(
+                        "Saving vocabulary to {}: vocabulary indices are not consecutive."
+                        " Please check that the vocabulary is not corrupted!".format(vocab_file)
+                    )
+                    index = token_index
+                writer.write(token + "\n")
+                index += 1
+        return (vocab_file,)
+
+
+class BasicTokenizer(object):
+    """
+    Constructs a BasicTokenizer that will run basic
+    tokenization (punctuation splitting, lower casing, etc.).
+    """
+
+    def __init__(self, do_lower_case=True, never_split=None, tokenize_chinese_chars=True):
+        """Constructs a BasicTokenizer.
+        Args:
+            **do_lower_case**: Whether to lower case the input.
+            **never_split**: (`optional`) list of str
+                Kept for backward compatibility purposes.
+                Now implemented directly at the base class level
+                (see :func:`PreTrainedTokenizer.tokenize`)
+                List of token not to split.
+            **tokenize_chinese_chars**: (`optional`) boolean (default True)
+                Whether to tokenize Chinese characters.
+                This should likely be deactivated for Japanese:
+                see: https://github.com/huggingface/pytorch-pretrained-BERT/issues/328
+        """
+        if never_split is None:
+            never_split = []
+        self.do_lower_case = do_lower_case
+        self.never_split = set(never_split)
+        self.tokenize_chinese_chars = tokenize_chinese_chars
+
+    def tokenize(self, text, never_split=None):
+        """
+        Basic Tokenization of a piece of text.
+        Split on "white spaces" only, for sub-word tokenization, see WordPieceTokenizer.
+
+        Args:
+            **never_split**: (`optional`) list of str
+                Kept for backward compatibility purposes.
+                Now implemented directly at the base class level
+                (see :func:`PreTrainedTokenizer.tokenize`)
+                List of token not to split.
+        """
+        # union() returns a new set by concatenating the two sets.
+        never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
+        text = self._clean_text(text)
+
+        # This was added on November 1st, 2018 for the multilingual and Chinese
+        # models. This is also applied to the English models now, but it doesn't
+        # matter since the English models were not trained on any Chinese data
+        # and generally don't have any Chinese data in them (there are Chinese
+        # characters in the vocabulary because Wikipedia does have some Chinese
+        # words in the English Wikipedia.).
+        if self.tokenize_chinese_chars:
+            text = self._tokenize_chinese_chars(text)
+        orig_tokens = whitespace_tokenize(text)
+        split_tokens = []
+        for token in orig_tokens:
+            if self.do_lower_case and token not in never_split:
+                token = token.lower()
+                token = self._run_strip_accents(token)
+            split_tokens.extend(self._run_split_on_punc(token, never_split))
+
+        output_tokens = whitespace_tokenize(" ".join(split_tokens))
+        return output_tokens
+
+    def _run_strip_accents(self, text):
+        """Strips accents from a piece of text."""
+        text = unicodedata.normalize("NFD", text)
+        output = []
+        for char in text:
+            cat = unicodedata.category(char)
+            if cat == "Mn":
+                continue
+            output.append(char)
+        return "".join(output)
+
+    def _run_split_on_punc(self, text, never_split=None):
+        """Splits punctuation on a piece of text."""
+        if never_split is not None and text in never_split:
+            return [text]
+        chars = list(text)
+        i = 0
+        start_new_word = True
+        output = []
+        while i < len(chars):
+            char = chars[i]
+            if _is_punctuation(char):
+                output.append([char])
+                start_new_word = True
+            else:
+                if start_new_word:
+                    output.append([])
+                start_new_word = False
+                output[-1].append(char)
+            i += 1
+
+        return ["".join(x) for x in output]
+
+    def _tokenize_chinese_chars(self, text):
+        """Adds whitespace around any CJK character."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if self._is_chinese_char(cp):
+                output.append(" ")
+                output.append(char)
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+    def _is_chinese_char(self, cp):
+        """Checks whether CP is the codepoint of a CJK character."""
+        # This defines a "chinese character" as anything in the CJK Unicode block:
+        #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+        #
+        # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+        # despite its name. The modern Korean Hangul alphabet is a different block,
+        # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+        # space-separated words, so they are not treated specially and handled
+        # like the all of the other languages.
+        if (
+            (cp >= 0x4E00 and cp <= 0x9FFF)
+            or (cp >= 0x3400 and cp <= 0x4DBF)  #
+            or (cp >= 0x20000 and cp <= 0x2A6DF)  #
+            or (cp >= 0x2A700 and cp <= 0x2B73F)  #
+            or (cp >= 0x2B740 and cp <= 0x2B81F)  #
+            or (cp >= 0x2B820 and cp <= 0x2CEAF)  #
+            or (cp >= 0xF900 and cp <= 0xFAFF)
+            or (cp >= 0x2F800 and cp <= 0x2FA1F)  #
+        ):  #
+            return True
+
+        return False
+
+    def _clean_text(self, text):
+        """Performs invalid character removal and whitespace cleanup on text."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if cp == 0 or cp == 0xFFFD or _is_control(char):
+                continue
+            if _is_whitespace(char):
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+
+class BasicTokenizerWithChineseWWM(BasicTokenizer):
+    """Pre-segmentation for Chinese sentences, which will be used in whole word mask."""
+
+    def __init__(self, do_lower_case=True, never_split=None, tokenize_chinese_chars=True):
+        super(BasicTokenizerWithChineseWWM, self).__init__(
+            do_lower_case=do_lower_case,
+            never_split=never_split,
+            tokenize_chinese_chars=tokenize_chinese_chars,
+        )
+        try:
+            import jieba
+
+            self.pre_tokenizer = lambda x: jieba.lcut(x, HMM=False)
+        except ImportError:
+            raise (ImportError("Chinese whole word mask need jieba"))
+
+    def _tokenize_chinese_chars(self, text):
+        """For Chinese pieces, uses jieba to segment the words and
+        adds whitespace around CJK character."""
+        output = []
+        piece = ""
+        for char in text:
+            cp = ord(char)
+            if self._is_chinese_char(cp):
+                piece += char
+            else:
+                chinese_words = self.pre_tokenizer(piece)
+                for word in chinese_words:
+                    output.append(" ")
+                    output.append(word)
+                    output.append(" ")
+                output.append(char)
+                piece = ""
+
+        chinese_words = self.pre_tokenizer(piece)
+        for word in chinese_words:
+            output.append(" ")
+            output.append(word)
+            output.append(" ")
+
+        return "".join(output)
+
+
+class WordpieceTokenizer(object):
+    """Runs WordPiece tokenization."""
+
+    def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
+        self.vocab = vocab
+        self.unk_token = unk_token
+        self.max_input_chars_per_word = max_input_chars_per_word
+
+    def tokenize(self, text):
+        """Tokenizes a piece of text into its word pieces.
+        This uses a greedy longest-match-first algorithm to perform tokenization
+        using the given vocabulary.
+        For example:
+          input = "unaffable"
+          output = ["un", "##aff", "##able"]
+
+          input = "有没有"
+          output = ["有", "##没", "##有"]
+        Args:
+          text: A single token or whitespace separated tokens. This should have
+            already been passed through `BasicTokenizer`.
+        Returns:
+          A list of wordpiece tokens.
+        """
+
+        output_tokens = []
+        for token in whitespace_tokenize(text):
+            chars = list(token)
+            if len(chars) > self.max_input_chars_per_word:
+                output_tokens.append(self.unk_token)
+                continue
+
+            is_bad = False
+            start = 0
+            sub_tokens = []
+            while start < len(chars):
+                end = len(chars)
+                cur_substr = None
+                while start < end:
+                    substr = "".join(chars[start:end])
+                    if start > 0:
+                        substr = "##" + substr
+
+                    if substr.startswith("##"):
+                        if _is_chinese_substr(substr):
+                            if substr[2:] in self.vocab:  # for Chinese substr
+                                cur_substr = substr
+                                break
+                        else:
+                            if substr in self.vocab:  # for English substr
+                                cur_substr = substr
+                                break
+                    else:
+                        if (
+                            substr in self.vocab
+                        ):  # non-substr, maybe character or whole Chinese word
+                            cur_substr = substr
+                            break
+                    end -= 1
+
+                if cur_substr is None:
+                    is_bad = True
+                    break
+
+                sub_tokens.append(cur_substr)
+                start = end
+
+            if is_bad:
+                output_tokens.append(self.unk_token)
+            else:
+                output_tokens.extend(sub_tokens)
+        return output_tokens

libai/tokenizer/tokenization_gpt2.py

+# coding=utf-8
+# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tokenization classes for OpenAI GPT (BPE)."""
+
+import json
+import logging
+import os
+from functools import lru_cache
+from io import open
+from typing import List, Optional
+
+import regex as re
+
+from .tokenization_base import PreTrainedTokenizer
+
+logger = logging.getLogger(__name__)
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "vocab.json",
+    "merges_file": "merges.txt",
+}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {"gpt2": "https://huggingface.co/gpt2/resolve/main/vocab.json"},
+    "merges_file": {"gpt2": "https://huggingface.co/gpt2/resolve/main/merges.txt"},
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "gpt2": 1024,
+}
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping
+    to whitespace/control characters the bpe code barfs on.
+
+    The reversible bpe codes work on unicode strings. This means you need a large # of unicode
+    characters in your vocab if you want to avoid UNKs. When you're at something like a 10B token
+    dataset you end up needing around 5K for decent coverage. This is a significant percentage
+    of your normal, say, 32K bpe vocab. To avoid that, we want lookup tables between utf-8
+    bytes and unicode strings.
+    """
+    bs = (
+        list(range(ord("!"), ord("~") + 1))
+        + list(range(ord("¡"), ord("¬") + 1))
+        + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    cs = bs[:]
+    n = 0
+    for b in range(2 ** 8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2 ** 8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """
+    Return set of symbol pairs in a word.
+
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+class GPT2Tokenizer(PreTrainedTokenizer):
+    """
+    Construct a GPT-2 tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    Args:
+        vocab_file (:obj:`str`):
+            Path to the vocabulary file.
+        merges_file (:obj:`str`):
+            Path to the merges file.
+        errors (:obj:`str`, `optional`, defaults to :obj:`"replace"`):
+            Paradigm to follow when decoding bytes to UTF-8. See `bytes.decode
+            <https://docs.python.org/3/library/stdtypes.html#bytes.decode>`__ for more information.
+        unk_token (:obj:`str`, `optional`, defaults to :obj:`<|endoftext|>`):
+            The unknown token. A token that is not in the vocabulary cannot be
+            converted to an ID and is set to be this token instead.
+        bos_token (:obj:`str`, `optional`, defaults to :obj:`<|endoftext|>`):
+            The beginning of sequence token.
+        eos_token (:obj:`str`, `optional`, defaults to :obj:`<|endoftext|>`):
+            The end of sequence token.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+
+    def __init__(
+        self,
+        vocab_file,
+        merges_file,
+        errors="replace",
+        unk_token="<|endoftext|>",
+        bos_token="<|endoftext|>",
+        eos_token="<|endoftext|>",
+        add_bos_token=False,
+        **kwargs,
+    ):
+        super(GPT2Tokenizer, self).__init__(
+            bos_token=bos_token, eos_token=eos_token, unk_token=unk_token, **kwargs
+        )
+
+        self.encoder = json.load(open(vocab_file, encoding="utf-8"))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.errors = errors  # how to handle errors in decoding
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        bpe_data = open(merges_file, encoding="utf-8").read().split("\n")[1:-1]
+        bpe_merges = [tuple(merge.split()) for merge in bpe_data]
+        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
+        self.cache = {}
+
+        # Should haved added re.IGNORECASE so BPE merges can happen for
+        # capitalized versions of contractions
+        self.pat = re.compile(
+            r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
+        )
+        self.add_bos_token = add_bos_token
+
+    @property
+    def vocab_size(self):
+        return len(self.encoder)
+
+    def get_vocab(self):
+        return dict(self.encoder, **self.added_tokens_encoder)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token
+
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:  # noqa
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = " ".join(word)
+        self.cache[token] = word
+        return word
+
+    def _tokenize(self, text):
+        """Tokenize a string."""
+        bpe_tokens = []
+        for token in re.findall(self.pat, text):
+            # Maps all our bytes to unicode strings, avoiding control tokens
+            # of the BPE (spaces in our case)
+            token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
+            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
+        return bpe_tokens
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.decoder.get(index)
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) to a single string."""
+        text = "".join(tokens)
+        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
+        return text
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """Add special tokens to a sequence or a pair of sequence.
+        GPT2 format sentence input:
+
+        - single sequence: <|endoftext|> tokens_a
+        - pair of sequences: <|endoftext|> tokens_a <|endoftext|> tokens_b
+
+        Args:
+            token_ids_0 (List[int]): The token ids of sentence 0.
+            token_ids_1 (List[int], optional): The token ids of sentence 1. Defaults to None.
+
+        Returns:
+            :obj:`List[str]`: The sequence after adding special toekens.
+        """
+        if self.add_bos_token:
+            bos = [self.bos_token_id]
+        else:
+            bos = []
+
+        if token_ids_1 is None:
+            return bos + token_ids_0
+
+        return bos + token_ids_0 + bos + token_ids_1
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"],
+        )
+        merge_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"],
+        )
+
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self.encoder, ensure_ascii=False))
+
+        index = 0
+        with open(merge_file, "w", encoding="utf-8") as writer:
+            writer.write("#version: 0.2\n")
+            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
+                if index != token_index:
+                    logger.warning(
+                        f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
+                        " Please check that the tokenizer is not corrupted!"
+                    )
+                    index = token_index
+                writer.write(" ".join(bpe_tokens) + "\n")
+                index += 1
+
+        return (vocab_file, merge_file)

libai/tokenizer/tokenization_roberta.py

+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tokenization classes for RoBERTa."""
+
+import json
+import logging
+import os
+from functools import lru_cache
+from typing import List, Optional, Tuple
+
+import regex as re
+
+from .tokenization_base import PreTrainedTokenizer
+
+logger = logging.getLogger(__name__)
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "vocab.json",
+    "merges_file": "merges.txt",
+}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "roberta-base": "https://huggingface.co/roberta-base/resolve/main/vocab.json",
+        "roberta-large": "https://huggingface.co/roberta-large/resolve/main/vocab.json",
+        "roberta-large-mnli": "https://huggingface.co/roberta-large-mnli/resolve/main/vocab.json",
+    },
+    "merges_file": {
+        "roberta-base": "https://huggingface.co/roberta-base/resolve/main/merges.txt",
+        "roberta-large": "https://huggingface.co/roberta-large/resolve/main/merges.txt",
+        "roberta-large-mnli": "https://huggingface.co/roberta-large-mnli/resolve/main/merges.txt",
+    },
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "roberta-base": 512,
+    "roberta-large": 512,
+    "roberta-large-mnli": 512,
+}
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to
+    whitespace/control characters the bpe code barfs on.
+
+    The reversible bpe codes work on unicode strings. This means you need a large # of unicode
+    characters in your vocab if you want to avoid UNKs. When you're at something like a 10B token
+    dataset you end up needing around 5K for decent coverage. This is a significant percentage of
+    your normal, say, 32K bpe vocab. To avoid that, we want lookup tables between utf-8 bytes and
+    unicode strings.
+    """
+    bs = (
+        list(range(ord("!"), ord("~") + 1))
+        + list(range(ord("¡"), ord("¬") + 1))
+        + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    cs = bs[:]
+    n = 0
+    for b in range(2 ** 8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2 ** 8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """
+    Return set of symbol pairs in a word.
+
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+class RobertaTokenizer(PreTrainedTokenizer):
+    """Constructs a RoBERTa tokenizer, derived from the GPT-2 tokenizer,
+    using byte-level Byte-Pair-Encoding.
+
+    Args:
+        vocab_file (:obj:`str`):
+            Path to the vocabulary file.
+        merges_file (:obj:`str`):
+            Path to the merges file.
+        errors (:obj:`str`, `optional`, defaults to :obj:`"replace"`):
+            Paradigm to follow when decoding bytes to UTF-8. See `bytes.decode
+            <https://docs.python.org/3/library/stdtypes.html#bytes.decode>`__ for more information.
+        bos_token (:obj:`str`, `optional`, defaults to `<s>`):
+            The beginning of sequence token.
+        eos_token (:obj:`str`, `optional`, defaults to `</s>`):
+            The end of sequence token.
+        cls_token (:obj:`str`, `optional`, defaults to `<s>`):
+            The first token of the sequence when built with special tokens.
+        unk_token (:obj:`str`, `optional`, defaults to `<unk>`):
+            The unknown token. A token that is not in the vocabulary cannot be
+            converted to an ID and is set to be this token instead.
+        pad_token (:obj:`str`, `optional`, defaults to `<pad>`): A special token
+            used to make arrays of tokens the same size for batching purpose.
+            Will then be ignored by attention mechanisms or loss computation.
+        mask_token (:obj:`str`, `optional`, defaults to `<mask>`): A special token
+            representing a masked token (used by masked-language modeling pretraining
+            objectives, like BERT).
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        merges_file,
+        errors="replace",
+        bos_token="<s>",
+        eos_token="</s>",
+        sep_token="</s>",
+        cls_token="<s>",
+        unk_token="<unk>",
+        pad_token="<pad>",
+        mask_token="<mask>",
+        add_bos_token=False,
+        **kwargs,
+    ):
+        super(RobertaTokenizer, self).__init__(
+            errors=errors,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            sep_token=sep_token,
+            cls_token=cls_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            **kwargs,
+        )
+
+        with open(vocab_file, encoding="utf-8") as file:
+            self.encoder = json.load(file)
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.errors = errors
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        with open(merges_file, encoding="utf-8") as file:
+            bpe_merges = file.read().split("\n")[1:-1]
+        bpe_merges = [tuple(merge.split()) for merge in bpe_merges]
+        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
+        self.cache = {}
+        self.pat = re.compile(
+            r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
+        )
+        self.add_bos_token = add_bos_token
+
+    @property
+    def vocab_size(self):
+        return len(self.encoder)
+
+    def get_vocab(self):
+        return dict(self.encoder, **self.added_tokens_encoder)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token
+
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+                else:
+                    new_word.extend(word[i:j])
+                    i = j
+
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = " ".join(word)
+        self.cache[token] = word
+        return word
+
+    def _tokenize(self, text):
+        """Tokenize a string."""
+        bpe_tokens = []
+        for token in re.findall(self.pat, text):
+            token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
+            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
+        return bpe_tokens
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.decoder.get(index)
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        text = "".join(tokens)
+        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
+        return text
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """Add special tokens to a sequence or a pair of sequence.
+        RoBERTa format sentence input:
+
+        - single sequence: [CLS] tokens_a [SEP]
+        - pair of sequences: [CLS] tokens_a [SEP] tokens_b [SEP]
+
+        Args:
+            token_ids_0 (List[int]): The token ids of sentence 0.
+            token_ids_1 (List[int], optional): The token ids of sentence 1. Defaults to None.
+
+        Returns:
+            :obj:`List[str]`: The sequence after adding special toekens.
+        """
+        if self.add_bos_token:
+            cls = [self.cls_token_id]
+            sep = [self.sep_token_id]
+        else:
+            cls = []
+            sep = []
+
+        if token_ids_1 is None:
+            return cls + token_ids_0 + sep
+
+        return cls + token_ids_0 + sep + token_ids_1 + sep
+
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"],
+        )
+        merge_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"],
+        )
+
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self.encoder, ensure_ascii=False))
+
+        index = 0
+        with open(merge_file, "w", encoding="utf-8") as writer:
+            writer.write("#version: 0.2\n")
+            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
+                if index != token_index:
+                    logger.warning(
+                        f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
+                        " Please check that the tokenizer is not corrupted!"
+                    )
+                    index = token_index
+                writer.write(" ".join(bpe_tokens) + "\n")
+                index += 1
+
+        return vocab_file, merge_file
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Create a mask from the two sequences passed to be used in a sequence-pair
+        classification task. RoBERTa does not make use of token type ids, therefore
+        a list of zeros is returned.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of zeros.
+        """
+        sep = [self.sep_token_id]
+        cls = [self.cls_token_id]
+
+        if token_ids_1 is None:
+            return len(cls + token_ids_0 + sep) * [0]
+        return len(cls + token_ids_0 + sep + sep + token_ids_1 + sep) * [0]

libai/tokenizer/tokenization_t5.py

+# coding=utf-8
+# Copyright 2018 T5 Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tokenization class for Google T5 (sentence piece)."""
+
+import logging
+import os
+import warnings
+from shutil import copyfile
+from typing import List, Optional
+
+import regex as re
+import sentencepiece as spm
+
+from .tokenization_base import PreTrainedTokenizer
+
+logger = logging.getLogger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {"t5-base": "https://huggingface.co/t5-base/resolve/main/spiece.model"}
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "t5-base": 512,
+}
+
+
+class T5Tokenizer(PreTrainedTokenizer):
+    """
+    Construct a T5 tokenizer. Based on `SentencePiece <https://github.com/google/sentencepiece>`.
+
+    Args:
+        vocab_file (:obj:`str`):
+            Path to the vocabulary file.
+        eos_token (:obj:`str`, `optional`, defaults to :obj:`"</s>"`):
+            The end of sequence token.
+        unk_token (:obj:`str`, `optional`, defaults to :obj:`"<unk>"`):
+            The unknown token. A token that is not in the vocabulary cannot
+            be converted to an ID and is set to be this token instead.
+        pad_token (:obj:`str`, `optional`, defaults to :obj:`"<pad>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        extra_ids (:obj:`int`, `optional`, defaults to 100):
+            Add a number of extra ids added to the end of the vocabulary for use
+            as sentinels. These tokens are accessible as "<extra_id_{%d}>" where
+            "{%d}" is a number between 0 and extra_ids-1. Extra tokens are indexed
+            from the end of the vocabulary up to beginning ("<extra_id_0>" is the
+            last token in the vocabulary like in T5 preprocessing see `here
+            <https://github.com/google-research/text-to-text-transfer-transformer/blob/9fd7b14a769417be33bc6c850f9598764913c833/t5/data/preprocessors.py#L2117>`__).
+        additional_special_tokens (:obj:`List[str]`, `optional`):
+            Additional special tokens used by the tokenizer.
+
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+
+    def __init__(
+        self,
+        vocab_file,
+        eos_token="</s>",
+        unk_token="<unk>",
+        pad_token="<pad>",
+        extra_ids=100,
+        additional_special_tokens=None,
+        add_bos_token=False,
+        **kwargs,
+    ):
+        # Add extra_ids to the special token list
+        if extra_ids > 0 and additional_special_tokens is None:
+            additional_special_tokens = [f"<extra_id_{i}>" for i in range(extra_ids)]
+        elif extra_ids > 0 and additional_special_tokens is not None:
+            extra_tokens = len(
+                set(filter(lambda x: bool("extra_id" in str(x)), additional_special_tokens))
+            )
+            if extra_tokens != extra_ids:
+                raise ValueError(
+                    f"Both extra_ids ({extra_ids}) and additional_special_tokens "
+                    f"({additional_special_tokens}) are privided to T5Tokenizer. "
+                    "In this case the additional_special_tokens must include the extra_ids tokens"
+                )
+
+        super().__init__(
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            additional_special_tokens=additional_special_tokens,
+            **kwargs,
+        )
+
+        self.vocab_file = vocab_file
+        self._extra_ids = extra_ids
+
+        self.sp_model = spm.SentencePieceProcessor()
+        self.sp_model.Load(vocab_file)
+        self.add_bos_token = add_bos_token
+
+    @property
+    def vocab_size(self):
+        return self.sp_model.get_piece_size() + self._extra_ids
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text):
+        """Tokenize a string."""
+        pieces = self.sp_model.encode(text, out_type=str)
+        return pieces
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        if token.startswith("<extra_id_"):
+            match = re.match(r"<extra_id_(\d+)>", token)
+            num = int(match.group(1))
+            return self.vocab_size - num - 1
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        if index < self.sp_model.get_piece_size():
+            token = self.sp_model.IdToPiece(index)
+        else:
+            token = f"<extra_id_{self.vocab_size - 1 - index}>"
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) to a single string."""
+        current_sub_tokens = []
+        out_string = ""
+        for token in tokens:
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                out_string += self.sp_model.decode_pieces(current_sub_tokens) + token + " "
+                current_sub_tokens = []
+            else:
+                current_sub_tokens.append(token)
+        out_string += self.sp_model.decode_pieces(current_sub_tokens)
+        return out_string.strip()
+
+    def _add_eos_if_not_present(self, token_ids):
+        if not self.add_bos_token:
+            return token_ids
+        if len(token_ids) > 0 and token_ids[-1] == self.eos_token_id:
+            warnings.warn("This sequence already has {self.eos_token}.")
+            return token_ids
+        else:
+            return token_ids + [self.eos_token_id]
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """Add special tokens to a sequence or a pair of sequence.
+        T5 format sentence input:
+
+        - single sequence: tokens_a </s>
+        - pair of sequences: tokens_a </s> tokens_b </s>
+
+        Args:
+            token_ids_0 (List[int]): The token ids of sentence 0.
+            token_ids_1 (List[int], optional): The token ids of sentence 1. Defaults to None.
+
+        Returns:
+            :obj:`List[str]`: The sequence after adding special toekens.
+        """
+        token_ids_0 = self._add_eos_if_not_present(token_ids_0)
+        if token_ids_1 is None:
+            return token_ids_0
+        else:
+            token_ids_1 = self._add_eos_if_not_present(token_ids_1)
+            return token_ids_0 + token_ids_1
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        """Save the tokenizer vocabulary to a directory or file."""
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"],
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+            logger.info(f"Copy vocab file to {out_vocab_file}")
+
+        return (out_vocab_file,)

libai/utils/__init__.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

libai/utils/checkpoint.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import logging
+import os
+import shutil
+from collections import defaultdict
+from typing import Any, Dict, Iterable, List, NamedTuple, Optional, Tuple
+
+import numpy as np
+import oneflow as flow
+from oneflow import nn
+from termcolor import colored
+
+import libai.utils.distributed as dist
+from libai.utils.file_io import HTTPURLHandler, PathManagerBase
+
+
+class _IncompatibleKeys(
+    NamedTuple(
+        # pyre-fixme[10]: Name `IncompatibleKeys` is used but not defined.
+        "IncompatibleKeys",
+        [
+            ("missing_keys", List[str]),
+            ("unexpected_keys", List[str]),
+            # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+            # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+            # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+            ("incorrect_shapes", List[Tuple]),
+        ],
+    )
+):
+    pass
+
+
+class Checkpointer(object):
+    """
+    A checkpointer that can save/load model as well as extra checkpointable
+    objects.
+    """
+
+    # NOTE: only support data_parallel for saving model
+    # TODO: save model: support model_parallel and pipeline parallel
+
+    def __init__(
+        self,
+        model: nn.Module,
+        save_dir: str = "",
+        *,
+        save_to_disk: bool = True,
+        **checkpointables: object,
+    ):
+        """
+        Args:
+            model (nn.Module): model.
+            save_dir (str): a directory to save and find checkpoints.
+            save_to_disk (bool): if True, save checkpoint to disk, otherwise
+                disable saving for this checkpointer.
+            checkpointables (object): any checkpointable objects, i.e., objects
+                that have the `state_dict()` and `load_state_dict()` method. For
+                example, it can be used like
+                `Checkpointer(model, "dir", optimizer=optimizer)`.
+        """
+        self.model = model
+        self.checkpointables = copy.copy(checkpointables)
+        self.logger = logging.getLogger(__name__)
+        self.save_dir = save_dir
+        self.save_to_disk = save_to_disk
+        # Default PathManager, support HTTP URLs
+        # A user may want to use a different project-specific PathManagerBase'
+        self.path_manager: PathManagerBase = PathManagerBase()
+        self.path_manager.register_handler(HTTPURLHandler())
+
+    def save(self, name: str, **kwargs: Dict[str, str]):
+        """
+        Dump model and checkpointables to a file.
+
+        Args:
+            name (str): name of the file.
+            kwargs (dict): extra arbitrary data to save.
+        """
+
+        data = {}
+        data["model"] = self.model.state_dict()
+        for key, obj in self.checkpointables.items():
+            data[key] = obj.state_dict()
+        data.update(kwargs)
+
+        basename = name
+        save_dir = os.path.join(self.save_dir, basename)
+        assert os.path.basename(save_dir) == basename, basename
+        if not self.path_manager.exists(save_dir):
+            self.path_manager.mkdirs(save_dir)
+        self.logger.info("Saving checkpoint to {}".format(save_dir))
+
+        for save_name in data:
+            if save_name == "iteration":
+                continue
+            save_file = os.path.join(save_dir, save_name)
+            # If directory existing, remove it for saving
+            if self.path_manager.exists(save_file):
+                self.path_manager.mkdirs(save_file)
+
+            flow.save(data[save_name], save_file, global_dst_rank=0)
+
+        if basename != "model_best":
+            self.tag_last_checkpoint(basename)
+
+    def load(self, path: str, checkpointables: Optional[List[str]] = None) -> object:
+        """
+        Load from the given checkpoint. When path points to network file, this
+        function has to be called on all ranks.
+
+        Args:
+            path (str): path or url to the checkpoint. If empty, will not load
+                anything.
+            checkpointables (list): List of checkpointable names to load. If not
+                specified (None), will load all the possible checkpointables.
+
+        Returns:
+            dict:
+                extra data loaded from the checkpoint that has not been
+                processed. For example, those saved with
+                :meth:`.save(**extra_data)`.
+        """
+        if not path:
+            # no checkpoint provided
+            self.logger.info("No checkpoint found. Training model from scratch")
+            return {}
+        self.logger.info("Loading checkpoint from {}".format(path))
+
+        checkpoint = self._load_file(path)
+        incompatible = self._load_model(checkpoint)
+        if incompatible is not None:  # handle some existing subclasses that returns None
+            self._log_incompatible_keys(incompatible)
+
+        for key in self.checkpointables if checkpointables is None else checkpointables:
+            if key in checkpoint:  # pyre-ignore
+                self.logger.info("Loading {} from {}".format(key, path))
+                obj = self.checkpointables[key]
+                obj.load_state_dict(checkpoint.pop(key))  # pyre-ignore
+
+        # return any further checkpoint data
+        return checkpoint
+
+    def has_checkpoint(self):
+        """
+        Returns:
+            bool: whether a checkpoint exists in the target directory.
+        """
+        save_file = os.path.join(self.save_dir, "last_checkpoint")
+        return self.path_manager.exists(save_file)
+
+    def get_checkpoint_file(self):
+        """
+        Returns:
+            str: The latest checkpoint file in target directory.
+        """
+        save_file = os.path.join(self.save_dir, "last_checkpoint")
+        try:
+            # load checkpoint file in rank0
+            if flow.env.get_rank() == 0:
+                with open(save_file, "r") as f:
+                    last_saved = f.read().strip()
+            else:
+                last_saved = None
+            # broadcast checkpoint file to other ranks
+            last_saved = dist.broadcast_py_object(last_saved, src=0)
+        except IOError:
+            # if file doesn't exist, maybe because it has just been
+            # deleted by a separate process
+            return ""
+        return os.path.join(self.save_dir, last_saved)
+
+    def resume_or_load(self, path: str, *, resume: bool = True):
+        """
+        If `resume` is True, this method attempts to resume from the last
+        checkpoint (if exists). Otherwise, load checkpoint from the given path.
+        This is useful when restarting an interrupted training job.
+        Args:
+            path (str): path to the checkpoint.
+            resume (bool): if True, resume from the last checkpoint if it exists.
+        Returns:
+            same as :meth:`load`.
+        """
+        if resume and self.has_checkpoint():
+            path = self.get_checkpoint_file()
+            return self.load(path)
+        else:
+            return self.load(path, checkpointables=[])
+
+    def tag_last_checkpoint(self, last_filename_basename: str):
+        """
+        Tag the last checkpoint.
+        Args:
+            last_filename_basename (str): the basename of the last filename.
+        """
+        save_file = os.path.join(self.save_dir, "last_checkpoint")
+        with self.path_manager.open(save_file, "w") as f:
+            f.write(last_filename_basename)  # pyre-ignore
+
+    def _load_file(self, f: str):
+        """
+        Load a checkpoint file. Can be overwritten by subclasses to support
+        different formats.
+        Args:
+            f (str): a locally mounted file path.
+        Returns:
+            dict: with keys "model" and optionally others that are saved by
+                the checkpointer dict["model"] must be a dict which maps strings
+                to flow.Tensor or numpy arrays.
+        """
+        data = {}
+        keys = self.path_manager.ls(f)
+        # broadcast checkpointer keys to other ranks
+        keys = dist.broadcast_py_object(keys, src=0)
+        for key in keys:
+            data[key] = flow.load(os.path.join(f, key), global_src_rank=0)
+        try:
+            data["iter"] = int(f.split("_")[-1])
+        except:  # noqa
+            self.logger.info(f"iter info in {f} not found, set iter to 0")
+            data["iter"] = 0
+        return data
+
+    def _load_model(self, checkpoint: Any):
+        """
+        Load weights from a checkpoint.
+        Args:
+            checkpoint (Any): checkpoint contains the weights.
+        """
+        checkpoint_state_dict = checkpoint.pop("model")
+        self._convert_ndarray_to_tensor(checkpoint_state_dict)
+
+        # if the state_dict comes from a model that was wrapped in a
+        # DataParallel or DistributedDataParallel during serialization,
+        # remove the "module" prefix before performing the matching.
+        _strip_prefix_if_present(checkpoint_state_dict, "module.")
+
+        model_state_dict = self.model.state_dict()
+        incorrect_shapes = []
+        for k in list(checkpoint_state_dict.keys()):
+            if k in model_state_dict:
+                shape_model = tuple(model_state_dict[k].shape)
+                shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
+                if shape_model != shape_checkpoint:
+                    incorrect_shapes.append((k, shape_checkpoint, shape_model))
+                    checkpoint_state_dict.pop(k)
+
+        incompatible = self.model.load_state_dict(checkpoint_state_dict, strict=False)
+        return _IncompatibleKeys(
+            missing_keys=incompatible.missing_keys,
+            unexpected_keys=incompatible.unexpected_keys,
+            incorrect_shapes=incorrect_shapes,
+        )
+
+    def _log_incompatible_keys(self, incompatible: _IncompatibleKeys) -> None:
+        """
+        Log information about the incompatible keys returned by ``_load_model``.
+        """
+        for k, shape_checkpoint, shape_model in incompatible.incorrect_shapes:
+            self.logger.warning(
+                "Skip loading parameter '{}' to the model due to incompatible "
+                "shapes: {} in the checkpoint but {} in the "
+                "model! You might want to double check if this is expected.".format(
+                    k, shape_checkpoint, shape_model
+                )
+            )
+        if incompatible.missing_keys:
+            missing_keys = _filter_reused_missing_keys(self.model, incompatible.missing_keys)
+            if missing_keys:
+                self.logger.info(get_missing_parameters_message(missing_keys))
+        if incompatible.unexpected_keys:
+            self.logger.info(get_unexpected_parameters_message(incompatible.unexpected_keys))
+
+    def _convert_ndarray_to_tensor(self, state_dict: dict):
+        """
+        In-place convert all numpy arrays in the state_dict to flow tensor.
+        Args:
+            state_dict (dict): a state-dict to be loaded to the model.
+        """
+        # model could be an OrderedDict with _metadata attribute
+        # (as returned by oneflow's state_dict()). We should preserve these
+        # properties.
+        for k in list(state_dict.keys()):
+            v = state_dict[k]
+            if not isinstance(v, np.ndarray) and not isinstance(v, flow.Tensor):
+                raise ValueError("Unsupported type found in checkpoint! {}: {}".format(k, type(v)))
+            # If it's local tensor, convert it to global tensor.
+            if not v.is_global:
+                if k in self.model.state_dict():
+                    model_v = self.model.state_dict()[k]
+                    state_dict[k] = v.to_global(sbp=model_v.sbp, placement=model_v.placement)
+
+
+class PeriodicCheckpointer:
+    """
+    Save checkpoints periodically. When `.step(iteration)` is called, it will
+    execute `checkpointer.save` on the given checkpointer, if iteration is a
+    multiple of period or if `max_iter` is reached.
+    """
+
+    def __init__(
+        self,
+        checkpointer: Checkpointer,
+        period: int,
+        max_iter: Optional[int] = None,
+        max_to_keep: Optional[int] = None,
+        file_prefix: str = "model",
+    ):
+        """
+        Args:
+            checkpointer (Any): the checkpointer object used to save
+            checkpoints.
+            period (int): the period to save checkpoint.
+            max_epoch (int): maximum number of epochs. When it is reached,
+                a checkpoint named "model_final" will be saved.
+        """
+        self.checkpointer = checkpointer
+        self.period = int(period)
+        self.max_iter = max_iter
+        if max_to_keep is not None:
+            assert max_to_keep > 0
+        self.max_to_keep = max_to_keep
+        self.recent_checkpoints: List[str] = []
+        self.file_prefix = file_prefix
+        self.path_manager: PathManagerBase = checkpointer.path_manager
+
+    def step(self, iteration: int, **kwargs: Any):
+        """
+        Perform the appropriate action at the given iteration.
+
+        Args:
+            iteration (int): the current epoch, ranged in [0, max_iter-1].
+            kwargs (Any): extra data to save, same as in
+                :meth:`Checkpointer.save`.
+        """
+        iteration = int(iteration)
+        additional_state = {"iteration": iteration}
+        additional_state.update(kwargs)
+
+        if (iteration + 1) % self.period == 0:
+            self.checkpointer.save(
+                "{}_{:07d}".format(self.file_prefix, iteration), **additional_state
+            )
+
+            if self.max_to_keep is not None:
+                self.recent_checkpoints.append(self.checkpointer.get_checkpoint_file())
+                if len(self.recent_checkpoints) > self.max_to_keep:
+                    file_to_delete = self.recent_checkpoints.pop(0)
+                    if (
+                        dist.is_main_process()
+                        and self.path_manager.exists(file_to_delete)
+                        and not file_to_delete.endswith(
+                            "{}_{:07d}".format(self.file_prefix, iteration)
+                        )
+                    ):
+                        if not self.path_manager.isfile(file_to_delete):
+                            shutil.rmtree(file_to_delete)
+                        else:
+                            self.path_manager.rm(file_to_delete)
+
+        if self.max_iter is not None:
+            if iteration >= self.max_iter - 1:
+                self.checkpointer.save(f"{self.file_prefix}_final", **additional_state)
+
+    def save(self, name: str, **kwargs: Any):
+        """
+        Same argument as :meth:`Checkpointer.save`.
+        Use this method to manually save checkpoints outside the schedule.
+
+        Args:
+            name (str): file name.
+            kwargs (Any): extra data to save, same as in
+                :meth:`Checkpointer.save`.
+        """
+        self.checkpointer.save(name, **kwargs)
+
+
+def _filter_reused_missing_keys(model: nn.Module, keys: List[str]) -> List[str]:
+    """
+    Filter "missing keys" to not include keys that have been loaded with another name.
+    """
+    keyset = set(keys)
+    param_to_names = defaultdict(set)  # param -> names that points to it
+    for module_prefix, module in _named_modules_with_dup(model):
+        for name, param in list(module.named_parameters(recurse=False)) + list(
+            module.named_buffers(recurse=False)  # pyre-ignore
+        ):
+            full_name = (module_prefix + "." if module_prefix else "") + name
+            param_to_names[param].add(full_name)
+    for names in param_to_names.values():
+        # if one name appears missing but its alias exists, then this
+        # name is not considered missing
+        if any(n in keyset for n in names) and not all(n in keyset for n in names):
+            [keyset.remove(n) for n in names if n in keyset]
+    return list(keyset)
+
+
+def get_missing_parameters_message(keys: List[str]) -> str:
+    """
+    Get a logging-friendly message to report parameter names (keys) that are in
+    the model but not found in a checkpoint.
+    Args:
+        keys (list[str]): List of keys that were not found in the checkpoint.
+    Returns:
+        str: message.
+    """
+    groups = _group_checkpoint_keys(keys)
+    msg = "Some model parameters or buffers are not found in the checkpoint:\n"
+    msg += "\n".join("  " + colored(k + _group_to_str(v), "blue") for k, v in groups.items())
+    return msg
+
+
+def get_unexpected_parameters_message(keys: List[str]) -> str:
+    """
+    Get a logging-friendly message to report parameter names (keys) that are in
+    the checkpoint but not found in the model.
+    Args:
+        keys (list[str]): List of keys that were not found in the model.
+    Returns:
+        str: message.
+    """
+    groups = _group_checkpoint_keys(keys)
+    msg = "The checkpoint state_dict contains keys that are not used by the model:\n"
+    msg += "\n".join("  " + colored(k + _group_to_str(v), "magenta") for k, v in groups.items())
+    return msg
+
+
+def _strip_prefix_if_present(state_dict: Dict[str, Any], prefix: str) -> None:
+    """
+    Strip the prefix in metadata, if any.
+    Args:
+        state_dict (OrderedDict): a state-dict to be loaded to the model.
+        prefix (str): prefix.
+    """
+    keys = sorted(state_dict.keys())
+    if not all(len(key) == 0 or key.startswith(prefix) for key in keys):
+        return
+
+    for key in keys:
+        newkey = key[len(prefix) :]
+        state_dict[newkey] = state_dict.pop(key)
+
+    # also strip the prefix in metadata, if any..
+    try:
+        metadata = state_dict._metadata  # pyre-ignore
+    except AttributeError:
+        pass
+    else:
+        for key in list(metadata.keys()):
+            # for the metadata dict, the key can be:
+            # '': for the DDP module, which we want to remove.
+            # 'module': for the actual model.
+            # 'module.xx.xx': for the rest.
+
+            if len(key) == 0:
+                continue
+            newkey = key[len(prefix) :]
+            metadata[newkey] = metadata.pop(key)
+
+
+def _group_checkpoint_keys(keys: List[str]) -> Dict[str, List[str]]:
+    """
+    Group keys based on common prefixes. A prefix is the string up to the final
+    "." in each key.
+    Args:
+        keys (list[str]): list of parameter names, i.e. keys in the model
+            checkpoint dict.
+    Returns:
+        dict[list]: keys with common prefixes are grouped into lists.
+    """
+    groups = defaultdict(list)
+    for key in keys:
+        pos = key.rfind(".")
+        if pos >= 0:
+            head, tail = key[:pos], [key[pos + 1 :]]
+        else:
+            head, tail = key, []
+        groups[head].extend(tail)
+    return groups
+
+
+def _group_to_str(group: List[str]) -> str:
+    """
+    Format a group of parameter name suffixes into a loggable string.
+    Args:
+        group (list[str]): list of parameter name suffixes.
+    Returns:
+        str: formated string.
+    """
+    if len(group) == 0:
+        return ""
+
+    if len(group) == 1:
+        return "." + group[0]
+
+    return ".{" + ", ".join(group) + "}"
+
+
+def _named_modules_with_dup(model: nn.Module, prefix: str = "") -> Iterable[Tuple[str, nn.Module]]:
+    """
+    The same as `model.named_modules()`, except that it includes
+    duplicated modules that have more than one name.
+    """
+    yield prefix, model
+    for name, module in model._modules.items():  # pyre-ignore
+        if module is None:
+            continue
+        submodule_prefix = prefix + ("." if prefix else "") + name
+        yield from _named_modules_with_dup(module, submodule_prefix)

libai/utils/distributed.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import dill
+import numpy as np
+import oneflow as flow
+from omegaconf import OmegaConf
+
+from libai.config import try_get_key
+
+logger = logging.getLogger(__name__)
+
+_DIST_UTIL = None
+
+
+def _merge_devices(devices):
+    num_gpus_per_node = get_world_size() // get_num_nodes()
+    node_devices = [node_id * num_gpus_per_node + device_id for node_id, device_id in devices]
+    return node_devices
+
+
+class _DistributeUtil(object):
+    def __init__(self, cfg):
+
+        self._init_distributed_env(cfg)
+        self._init_parallel_size(cfg)
+        self._init_placement_group(cfg)
+        self._init_parallel_hierarchy()
+
+    def _init_distributed_env(self, cfg):
+        """Initialize the distributed environment."""
+
+        num_nodes = get_num_nodes()
+        num_gpus_per_node = get_world_size() // num_nodes
+
+        if try_get_key(cfg, "num_gpus_per_node", default=num_gpus_per_node) != num_gpus_per_node:
+            # This means key(num_gpus_per_node) saved in config is not equal
+            # to environment variable.
+            # Give user a warning about inconsistent reproduce environment.
+            logger.warning(
+                "'train.dist.num_gpus_per_node' are not equal to environment variable. "
+                f"{cfg.num_gpus_per_node} != {num_gpus_per_node}"
+            )
+
+        if try_get_key(cfg, "num_nodes", default=num_nodes) != num_nodes:
+            logger.warning(
+                "'train.dist.num_nodes' are not equal to"
+                f"environment variable. {cfg.num_nodes} != {num_nodes}"
+            )
+
+        # Set the actual value to config
+        cfg.num_nodes = num_nodes
+        cfg.num_gpus_per_node = num_gpus_per_node
+
+        self._num_nodes = num_nodes
+        self._num_gpus_per_node = num_gpus_per_node
+        self._world_size = num_gpus_per_node * num_nodes
+
+        # Add set device type
+        self._device_type = try_get_key(cfg, "device_type", default="cuda")
+
+    def _init_parallel_size(self, cfg):
+
+        # tensor parallel size
+        self._tensor_parallel_size = min(cfg.tensor_parallel_size, self.world_size)
+        assert self.world_size % self._tensor_parallel_size == 0, (
+            f"world size ({self.world_size}) is not divisible by"
+            f" tensor parallel size ({self._tensor_parallel_size})"
+        )
+        # Set the actual tensor parallel size to cfg
+        cfg.tensor_parallel_size = self._tensor_parallel_size
+
+        # pipeline parallel size
+        self._pipeline_parallel_size = min(
+            cfg.pipeline_parallel_size, self.world_size // cfg.tensor_parallel_size
+        )
+        # Set the actual pipeline parallel size to cfg
+        cfg.pipeline_parallel_size = self._pipeline_parallel_size
+
+        if cfg.pipeline_parallel_size > 1:
+            assert (
+                try_get_key(cfg, "pipeline_num_layers") is not None
+            ), "cfg.train.dist.pipeline_num_layers must be set when run pipeline parallel"
+
+            assert cfg.pipeline_num_layers >= self._pipeline_parallel_size, (
+                f"number of layers ({cfg.pipeline_num_layers}) is less than"
+                f" pipeline model parallel size ({self._pipeline_parallel_size})"
+            )
+            if try_get_key(cfg, "custom_pipeline_stage_id") is not None:
+                assert OmegaConf.is_list(
+                    cfg.custom_pipeline_stage_id
+                ), "type of cfg.train.dist.custom_pipeline_stage_id must be list"
+                cfg.custom_pipeline_stage_id = list(cfg.custom_pipeline_stage_id)
+                assert max(cfg.custom_pipeline_stage_id) < self._world_size, (
+                    f"the element {max(cfg.custom_pipeline_stage_id)} in"
+                    " cfg.train.dist.custom_pipeline_stage_id is out of range"
+                    f" for total rank {self._world_size}"
+                )
+                assert len(cfg.custom_pipeline_stage_id) == cfg.pipeline_num_layers, (
+                    "the length of cfg.train.dist.custom_pipeline_stage_id"
+                    f" {len(cfg.custom_pipeline_stage_id)} must be equal to"
+                    " cfg.train.dist.pipeline_num_layers"
+                    f" {cfg.train.dist.pipeline_num_layers}"
+                )
+        else:
+            # no pipeline parallel, just set 10000
+            if try_get_key(cfg, "pipeline_num_layers") is None:
+                cfg.pipeline_num_layers = 10000
+
+        self._model_parallel_size = self._pipeline_parallel_size * self._tensor_parallel_size
+
+        assert self.world_size % self._model_parallel_size == 0, (
+            f"world size ({self.world_size}) is not divisible by"
+            f" tensor model parallel size ({self._tensor_parallel_size}) times"
+            f" pipeline model parallel size ({self._pipeline_parallel_size})"
+        )
+
+        # data parallel size
+        self._data_parallel_size = self.world_size // self._model_parallel_size
+        # Set the actual data parallel size to cfg
+        cfg.data_parallel_size = self._data_parallel_size
+
+    def _init_placement_group(self, cfg):
+        node_ids = [i // self.num_gpus_per_node for i in range(self.world_size)]
+        device_ids = list(range(self.num_gpus_per_node)) * self.num_nodes
+
+        # [(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3)]
+        devices = [(n, d) for n, d in zip(node_ids, device_ids)]
+        num_devices_per_stage = self.world_size // self._pipeline_parallel_size
+        stages_devices = [
+            _merge_devices(devices[i : (i + num_devices_per_stage)])
+            for i in range(0, self.world_size, num_devices_per_stage)
+        ]
+
+        # change pipeline_num_layers to make the middle stages contain more layers
+        if (
+            self._pipeline_parallel_size >= 4
+            and cfg.pipeline_num_layers >= 8
+            and cfg.pipeline_num_layers % self._pipeline_parallel_size == 0
+        ):
+            temp_num_layers_per_stage = cfg.pipeline_num_layers // self._pipeline_parallel_size
+            actual_pipeline_num_layers = cfg.pipeline_num_layers + min(
+                self._pipeline_parallel_size - 1, temp_num_layers_per_stage
+            )
+        else:
+            actual_pipeline_num_layers = cfg.pipeline_num_layers
+
+        num_layers_per_stage = actual_pipeline_num_layers // self._pipeline_parallel_size
+        stage_offset = actual_pipeline_num_layers % self._pipeline_parallel_size
+
+        # stage_offset can make the later stages contain more layers when pipeline_num_layers
+        # cannot be divided by pipeline_parallel_size.
+        # This can make pipeline parallel more memory efficient.
+        self._layer_stage_ids = []
+        for i in range(0, actual_pipeline_num_layers - stage_offset, num_layers_per_stage):
+            stage_id = i // num_layers_per_stage
+            if stage_id >= (self._pipeline_parallel_size - stage_offset):
+                self._layer_stage_ids.append(stage_id)
+            self._layer_stage_ids.extend([stage_id] * num_layers_per_stage)
+        self._layer_stage_ids = self._layer_stage_ids[: cfg.pipeline_num_layers]
+        # when pipeline_parallel_size > 1, we add pipeline_stage_id infomation into cfg
+        if cfg.pipeline_parallel_size > 1:
+            cfg.auto_pipeline_stage_id = self._layer_stage_ids
+            # set pipeline_stage_id by users' setting
+            if try_get_key(cfg, "custom_pipeline_stage_id") is not None:
+                self._layer_stage_ids = cfg.custom_pipeline_stage_id
+            cfg.actual_pipeline_stage_id = self._layer_stage_ids
+
+        self._layer_ranks = [stages_devices[stage_id] for stage_id in self._layer_stage_ids]
+
+    def _init_parallel_hierarchy(self):
+        if self.is_data_model_parallel():
+            self._parallel_hierarchy = (
+                self._data_parallel_size,
+                self._tensor_parallel_size,
+            )
+        else:
+            self._parallel_hierarchy = None
+
+    @property
+    def num_nodes(self):
+        return self._num_nodes
+
+    @property
+    def num_gpus_per_node(self):
+        return self._num_gpus_per_node
+
+    @property
+    def world_size(self):
+        return self._world_size
+
+    @property
+    def parallel_hierarchy(self):
+        return self._parallel_hierarchy
+
+    @property
+    def tensor_parallel_size(self):
+        return self._tensor_parallel_size
+
+    @property
+    def pipeline_parallel_size(self):
+        return self._pipeline_parallel_size
+
+    @property
+    def model_parallel_size(self):
+        return self._tensor_parallel_size
+
+    @property
+    def data_parallel_size(self):
+        return self._data_parallel_size
+
+    @property
+    def device_type(self):
+        return self._device_type
+
+    def set_device_type(self, device_type):
+        assert device_type in ["cpu", "cuda"], f"not supported for {device_type}"
+        self._device_type = device_type
+
+    def get_layer_ranks(self, layer_idx):
+        layer_ranks = self._layer_ranks[layer_idx]
+        if self._parallel_hierarchy is None:
+            return layer_ranks
+        else:
+            assert len(self._parallel_hierarchy) == 2
+            return np.asarray(layer_ranks).reshape(self._parallel_hierarchy).tolist()
+
+    def get_layer_stage_id(self, layer_idx):
+        return self._layer_stage_ids[layer_idx]
+
+    def is_tensor_model_parallel(self):
+        return self._tensor_parallel_size > 1
+
+    def is_data_parallel(self):
+        return self._data_parallel_size > 1
+
+    def is_pipeline_model_parallel(self):
+        return self._pipeline_parallel_size > 1
+
+    def is_data_model_parallel(self):
+        return self.is_tensor_model_parallel() and self.is_data_parallel()
+
+
+def setup_dist_util(cfg):
+    """Initialize the distributed environment with configuration.
+
+    Example:
+
+    .. code-block:: python
+
+        from omegaconf import DictConfig
+
+        # set the hybrid parallel distributed environment with 2D mesh GPUs
+        setup_dist_util(
+            DictConfig(
+                dict(
+                    data_parallel_size=2,
+                    tensor_parallel_size=2,
+                    pipeline_parallel_size=1,
+                )
+            )
+        )
+
+    """
+    global _DIST_UTIL
+    _DIST_UTIL = _DistributeUtil(cfg)
+
+
+def get_dist_util():
+    """Get distributed utils if it's been setup. Otherwise, initialize it with
+    single node/single gpu environment."""
+    global _DIST_UTIL
+    if _DIST_UTIL is None:
+        logger.warning(
+            "Distributed env is not set up, configure it by default (single node, single gpu)."
+        )
+        from omegaconf import DictConfig
+
+        setup_dist_util(
+            DictConfig(
+                dict(
+                    data_parallel_size=1,
+                    tensor_parallel_size=1,
+                    pipeline_parallel_size=1,
+                )
+            )
+        )
+    return _DIST_UTIL
+
+
+def get_layer_placement(layer_idx, device_type=None):
+    """
+    Get ``flow.placement`` object with the initialized distributed environment
+    according to the ``layer_idx``.
+
+    Args:
+        layer_idx (int): layer index indicating the rank groups. This is very useful for pipeline
+            parallelism training where different layers are on different ranks.
+        device_type (str, optional): device type. Defaults to "cuda".
+    """
+    dist_util = get_dist_util()
+    device_type = dist_util.device_type if device_type is None else device_type
+    if not flow.cuda.is_available() and device_type == "cuda":
+        device_type = "cpu"
+    return flow.placement(
+        device_type,
+        dist_util.get_layer_ranks(layer_idx),
+    )
+
+
+def get_nd_sbp(sbp_list):
+    """Get nd sbp signature list, which is consistent with 1D/2D mesh GPUs.
+
+    Args:
+        sbp_list (list): a sbp list with 2D mesh.
+
+    Returns:
+        A modified sbp list according to the initialized distributed environment.
+    """
+    assert isinstance(sbp_list, list)
+    assert len(sbp_list) == 2
+    assert all(isinstance(sbp, flow.sbp.sbp) for sbp in sbp_list)
+
+    dist_util = get_dist_util()
+    if dist_util.is_data_model_parallel():
+        return sbp_list
+    elif dist_util.is_data_parallel():
+        return sbp_list[:1]
+    elif dist_util.is_tensor_model_parallel():
+        return sbp_list[1:]
+    else:
+        return [flow.sbp.broadcast]
+
+
+def get_hidden_sbp():
+    """Hidden states sbp."""
+    return get_nd_sbp([flow.sbp.split(0), flow.sbp.broadcast])
+
+
+def get_data_parallel_rank():
+    dist_util = get_dist_util()
+    return (flow.env.get_rank() // dist_util.model_parallel_size) % dist_util.data_parallel_size
+
+
+def get_data_parallel_size():
+    dist_util = get_dist_util()
+    return dist_util.data_parallel_size
+
+
+def get_tensor_parallel_size():
+    dist_util = get_dist_util()
+    return dist_util.tensor_parallel_size
+
+
+def get_pipeline_parallel_size():
+    dist_util = get_dist_util()
+    return dist_util.pipeline_parallel_size
+
+
+def same_sbp(lhs_sbp, rhs_sbp):
+    """Determine if two sbp signatures are the same."""
+    assert len(lhs_sbp) == len(rhs_sbp)
+
+    for i in range(len(lhs_sbp)):
+        if lhs_sbp[i] != rhs_sbp[i]:
+            return False
+    return True
+
+
+def get_rank() -> int:
+    return flow.env.get_rank()
+
+
+def get_local_rank() -> int:
+    return flow.env.get_local_rank()
+
+
+def is_main_process() -> bool:
+    return get_rank() == 0
+
+
+def is_last_process() -> bool:
+    return get_rank() == get_world_size() - 1
+
+
+def get_world_size():
+    return flow.env.get_world_size()
+
+
+def get_num_nodes():
+    return flow.env.get_node_size()
+
+
+def set_device_type(device_type):
+    dist_util = get_dist_util()
+    dist_util.set_device_type(device_type)
+
+
+def broadcast_py_object(obj, src: int = 0):
+    rank = flow.env.get_rank()
+    if src == rank:
+        obj_bytes = dill.dumps(obj)
+        return dill.loads(flow._oneflow_internal.cpu_broadcast(obj_bytes, src))
+    else:
+        return dill.loads(flow._oneflow_internal.cpu_broadcast(None, src))
+
+
+def convert_to_distributed_default_setting(t):
+    """
+    Helper function to convert all eager local tensor in :attr:`nn.Module` in the model to
+    global tensor with data parallelism as default.
+    """
+    if not t.is_global:
+        return t.to_global(
+            sbp=get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+            placement=get_layer_placement(0),
+        )
+    else:
+        dist_util = get_dist_util()
+        device_type = dist_util.device_type
+        return t.to_global(placement=flow.placement(device_type, ranks=t.placement.ranks))
+
+
+def ttol(tensor, pure_local=False, ranks=None):
+    """Global tensor to local tensor."""
+    if tensor.is_global:
+        placement = tensor.placement if not ranks else flow.placement("cuda", ranks)
+        if pure_local:
+            tensor = tensor.to_global(placement=placement).to_local()
+        else:
+            tensor = tensor.to_global(
+                sbp=get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]), placement=placement
+            ).to_local()
+
+    return tensor
+
+
+def tton(tensor, local_only=False, ranks=None):
+    """Global tensor to numpy ndarray."""
+    if tensor.is_global:
+        tensor = ttol(tensor, local_only, ranks)
+
+    return tensor.numpy()
+
+
+def tensor_to_rank0(tensor, device="cuda", to_local=False):
+    """Global tensor to rank0."""
+    assert device in ["cpu", "cuda"], f"not supported for device:{device}"
+    if tensor.is_global:
+        # Consider if it's 2d mesh, ranks should be [[0]] instead of [0]
+        placement = flow.placement(device, ranks=[0] if tensor.placement.ranks.ndim == 1 else [[0]])
+        tensor = tensor.to_global(
+            sbp=get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]), placement=placement
+        )
+        if to_local:
+            tensor = ttol(tensor)
+    return tensor
+
+
+def synchronize():
+    """
+    Helper function to synchronize (barrier) among all processes when
+    using distributed training.
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return
+
+    flow.comm.barrier()

libai/utils/download.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import os
+import shutil
+from typing import Callable, List, Optional
+from urllib import request
+
+# --------------------------------------------------------
+# References:
+# https://github.com/facebookresearch/iopath/blob/main/iopath/common/download.py
+# --------------------------------------------------------
+
+
+def download(url: str, dir: str, *, filename: Optional[str] = None, progress: bool = True) -> str:
+    """
+    Download a file from a given URL to a directory. If file exists, will not
+        overwrite the existing file.
+    Args:
+        url (str):
+        dir (str): the directory to download the file
+        filename (str or None): the basename to save the file.
+            Will use the name in the URL if not given.
+        progress (bool): whether to use tqdm to draw a progress bar.
+    Returns:
+        str: the path to the downloaded file or the existing one.
+    """
+    os.makedirs(dir, exist_ok=True)
+    if filename is None:
+        filename = url.split("/")[-1]
+        assert len(filename), "Cannot obtain filename from url {}".format(url)
+    fpath = os.path.join(dir, filename)
+    logger = logging.getLogger(__name__)
+
+    if os.path.isfile(fpath):
+        logger.info("File {} exists! Skipping download.".format(filename))
+        return fpath
+
+    tmp = fpath + ".tmp"  # download to a tmp file first, to be more atomic.
+    try:
+        logger.info("Downloading from {} ...".format(url))
+        if progress:
+            import tqdm
+
+            def hook(t: tqdm.tqdm) -> Callable[[int, int, Optional[int]], None]:
+                last_b: List[int] = [0]
+
+                def inner(b: int, bsize: int, tsize: Optional[int] = None) -> None:
+                    if tsize is not None:
+                        t.total = tsize
+                    t.update((b - last_b[0]) * bsize)  # type: ignore
+                    last_b[0] = b
+
+                return inner
+
+            with tqdm.tqdm(  # type: ignore
+                unit="B", unit_scale=True, miniters=1, desc=filename, leave=True
+            ) as t:
+                tmp, _ = request.urlretrieve(url, filename=tmp, reporthook=hook(t))
+
+        else:
+            tmp, _ = request.urlretrieve(url, filename=tmp)
+        statinfo = os.stat(tmp)
+        size = statinfo.st_size
+        if size == 0:
+            raise IOError("Downloaded an empty file from {}!".format(url))
+        # download to tmp first and move to fpath, to make this function more
+        # atomic.
+        shutil.move(tmp, fpath)
+    except IOError:
+        logger.error("Failed to download {}".format(url))
+        raise
+    finally:
+        try:
+            os.unlink(tmp)
+        except IOError:
+            pass
+
+    logger.info("Successfully downloaded " + fpath + ". " + str(size) + " bytes.")
+    return fpath

libai/utils/events.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import datetime
+import json
+import logging
+import os
+import time
+from collections import defaultdict
+from contextlib import contextmanager
+
+from libai.utils.file_io import PathManager
+from libai.utils.history_buffer import HistoryBuffer
+
+# --------------------------------------------------------
+# References:
+# https://github.com/facebookresearch/detectron2/blob/main/detectron2/utils/events.py
+# --------------------------------------------------------
+
+
+__all__ = [
+    "get_event_storage",
+    "JSONWriter",
+    "CommonMetricPrinter",
+    "EventStorage",
+]
+
+_CURRENT_STORAGE_STACK = []
+
+
+def get_event_storage():
+    """
+    Returns:
+        The :class:`EventStorage` object that's currently being used.
+        Throw an error if no :class:`EventStorage` is currently enabled.
+    """
+    assert len(
+        _CURRENT_STORAGE_STACK
+    ), "get_event_storage() has to be called inside a 'with EventStorage(...)' context!"
+    return _CURRENT_STORAGE_STACK[-1]
+
+
+class EventWriter:
+    """
+    Base class for writers that obtain events from :class:`EventStorage` and process them.
+    """
+
+    def write(self):
+        raise NotImplementedError
+
+    def close(self):
+        pass
+
+
+class JSONWriter(EventWriter):
+    """
+    Write scalars to a json file.
+    It saves scalars as one json per line (instead of a big json) for easy parsing.
+    Example of parsing such a json file:
+
+    ::
+
+        $ cat metrics.json | jq -s '.[0:2]'
+        [
+          {
+            "data_time": 0.008433341979980469,
+            "iteration": 19,
+            "total_loss": 1.9228371381759644,
+            "lr": 0.007173333333333333,
+            "time": 0.25401854515075684
+          },
+          {
+            "data_time": 0.007216215133666992,
+            "iteration": 39,
+            "total_loss": 1.282649278640747,
+            "lr": 0.007706666666666667,
+            "time": 0.2490077018737793
+          }
+        ]
+        $ cat metrics.json | jq '.loss_mask'
+        0.7126231789588928
+        0.689423680305481
+        0.6776131987571716
+        ...
+    """
+
+    def __init__(self, json_file, window_size=20):
+        """
+        Args:
+            json_file (str): path to the json file. New data will be appended if the file exists.
+            window_size (int): the window size of median smoothing for the scalars whose
+                `smoothing_hint` are True.
+        """
+        self._file_handle = PathManager.open(json_file, "a")
+        self._window_size = window_size
+        self._last_write = -1
+
+    def write(self):
+        storage = get_event_storage()
+        to_save = defaultdict(dict)
+
+        for k, (v, iter) in storage.latest_with_smoothing_hint(self._window_size).items():
+            # keep scalars that have not been written
+            if iter <= self._last_write:
+                continue
+            to_save[iter][k] = v
+        if len(to_save):
+            all_iters = sorted(to_save.keys())
+            self._last_write = max(all_iters)
+
+        for itr, scalars_per_iter in to_save.items():
+            scalars_per_iter["iteration"] = itr
+            self._file_handle.write(json.dumps(scalars_per_iter, sort_keys=True) + "\n")
+        self._file_handle.flush()
+        try:
+            os.fsync(self._file_handle.fileno())
+        except AttributeError:
+            pass
+
+    def close(self):
+        self._file_handle.close()
+
+
+class TensorboardXWriter(EventWriter):
+    """
+    Write all scalars to a tensorboard file
+    """
+
+    def __init__(self, log_dir: str, window_size: int = 20, **kwargs):
+        """
+        Args:
+            log_dir (str): the directory to save the output events
+            window_size (int): the scalars will be median-smoothed by this window size
+
+            kwargs: other arguments passed to `tensorboardX.SummaryWriter(...)`
+        """
+        self._window_size = window_size
+        from tensorboardX import SummaryWriter
+
+        self._writer = SummaryWriter(log_dir=log_dir, **kwargs)
+        self._last_write = -1
+
+    def write(self):
+        storage = get_event_storage()
+        new_last_write = self._last_write
+        for k, (v, iter) in storage.latest_with_smoothing_hint(self._window_size).items():
+            if iter > self._last_write:
+                self._writer.add_scalar(k, v, iter)
+                new_last_write = max(new_last_write, iter)
+        self._last_write = new_last_write
+
+        # TODO: add write image
+
+        if len(storage._histograms) >= 1:
+            for params in storage._histograms:
+                self._writer.add_histogram_raw(**params)
+            storage.clear_histograms()
+
+    def close(self):
+        if hasattr(self, "_writer"):  # doesn't exist when the code fails at import
+            self._writer.close()
+
+
+class CommonMetricPrinter(EventWriter):
+    """
+    Print **common** metrics to the terminal, including
+    iteration time, ETA, memory, all losses, and the learning rate.
+    It also applies smoothing using a window of 20 elements.
+    It's meant to print common metrics in common ways.
+    To print something in more customized ways, please implement a similar printer by yourself.
+    """
+
+    def __init__(self, batch_size, max_iter):
+        """
+        Args:
+            max_iter (int): the maximum number of iterations to train.
+                Used to compute ETA.
+        """
+        self.logger = logging.getLogger(__name__)
+        self._batch_size = batch_size
+        self._max_iter = max_iter
+        self._last_write = None
+
+    def write(self):
+        storage = get_event_storage()
+        iteration = storage.iter
+        consumed_samples = storage.samples
+        if iteration == self._max_iter:
+            # This hook only reports training progress (loss, ETA, etc) but not other data,
+            # therefore do not write anything after training succeeds, even if this method
+            # is called.
+            return
+
+        try:
+            data_time = storage.history("data_time").avg(20)
+        except KeyError:
+            # they may not exist in the first few iterations (due to warmup)
+            # or when SimpleTrainer is not used
+            data_time = None
+
+        eta_string = None
+        try:
+            iter_time = storage.history("time").global_avg()
+            eta_seconds = storage.history("time").median(1000) * (self._max_iter - iteration - 1)
+            storage.put_scalar("eta_seconds", eta_seconds, smoothing_hint=False)
+            eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+        except KeyError:
+            iter_time = None
+            # estimate eta on our own - more noisy
+            if self._last_write is not None:
+                estimate_iter_time = (time.perf_counter() - self._last_write[1]) / (
+                    iteration - self._last_write[0]
+                )
+                eta_seconds = estimate_iter_time * (self._max_iter - iteration - 1)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+            self._last_write = (iteration, time.perf_counter())
+
+        try:
+            lr = "{:.2e}".format(storage.history("lr").latest())
+        except KeyError:
+            lr = "N/A"
+
+        max_mem_mb = None
+
+        # NOTE: max_mem is parsed by grep in "dev/parse_results.sh"
+        self.logger.info(
+            " {eta}{iter}  {sample}  {losses}  {time}{data_time} {tpt} lr: {lr}  {memory}".format(
+                eta=f"eta: {eta_string}  " if eta_string else "",
+                iter=f"iteration: {iteration}/{self._max_iter}",
+                sample=f"consumed_samples: {consumed_samples}",
+                losses="  ".join(
+                    [
+                        "{}: {:.4g}".format(k, v.median(200))
+                        for k, v in storage.histories().items()
+                        if "loss" in k
+                    ]
+                ),
+                time="time: {:.4f} s/iter  ".format(iter_time) if iter_time is not None else "",
+                data_time="data_time: {:.4f} s/iter".format(data_time)
+                if data_time is not None
+                else "",
+                tpt="total_throughput: {:.2f} samples/s".format(self._batch_size / iter_time)
+                if iter_time is not None
+                else "",
+                lr=lr,
+                memory="max_mem: {:.0f}M".format(max_mem_mb) if max_mem_mb is not None else "",
+            )
+        )
+
+
+class EventStorage:
+    """
+    The user-facing class that provides metric storage functionalities.
+    In the future we may add support for storing / logging other types of data if needed.
+    """
+
+    def __init__(self, start_iter=0):
+        """
+
+        Args:
+            start_iter (int): the iteration number to start with
+        """
+        self._history = defaultdict(HistoryBuffer)
+        self._smoothing_hints = {}
+        self._latest_scalars = {}
+        self._iter = start_iter
+        self._batch_size = 0
+        self._current_prefix = ""
+        self._vis_data = []
+        self._histograms = []
+
+    def put_image(self, img_name, img_tensor):
+        """
+        Add an `img_tensor` associated with `img_name` to be shown on
+        tensorboard.
+
+        Args:
+            img_name (str): The name of the image to put into tensorboard.
+            img_tensor (flow.Tensor or numpy.array): An `uint8` or `float`
+                Tensor of shape `[channel, height, width]` where `channel` is
+                3. The image format should be RGB. The elements in img_tensor
+                can either have values in [0, 1] (float32) or [0, 255] (uint8).
+                The `img_tensor` will be visualized in tensorboard.
+        """
+        self._vis_data.append((img_name, img_tensor, self._iter))
+
+    def put_scalar(self, name, value, smoothing_hint=True):
+        """
+        Add a scalar `value` to the `HistoryBuffer` associated with `name`.
+
+        Args:
+            smoothing_hint (bool): a 'hint' on whether this scalar is noisy and should be
+                smoothed when logged. The hint will be accessible through
+                :meth:`EventStorage.smoothing_hints`.  A writer may ignore the hint
+                and apply custom smoothing rule.
+                It defaults to True because most scalars we save need to be smoothed to
+                provide any useful signal.
+        """
+        name = self._current_prefix + name
+        history = self._history[name]
+        value = float(value)
+        history.update(value, self._iter)
+        self._latest_scalars[name] = (value, self._iter)
+
+        existing_hint = self._smoothing_hints.get(name)
+        if existing_hint is not None:
+            assert (
+                existing_hint == smoothing_hint
+            ), "Scalar {} was put with a different smoothing_hint!".format(name)
+        else:
+            self._smoothing_hints[name] = smoothing_hint
+
+    def put_scalars(self, *, smoothing_hint=True, **kwargs):
+        """
+        Put multiple scalars from keyword arguments.
+
+        Example:
+
+        .. code-block:: python
+
+            storage.put_scalars(loss=my_loss, accuracy=my_accuracy, smoothing_hint=True)
+        """
+        for k, v in kwargs.items():
+            self.put_scalar(k, v, smoothing_hint=smoothing_hint)
+
+    def history(self, name):
+        """
+        Returns:
+            HistoryBuffer: the scalar history for name
+        """
+        ret = self._history.get(name, None)
+        if ret is None:
+            raise KeyError("No history metric available for {}!".format(name))
+        return ret
+
+    def histories(self):
+        """
+        Returns:
+            dict[name -> HistoryBuffer]: the HistoryBuffer for all scalars
+        """
+        return self._history
+
+    def latest(self):
+        """
+        Returns:
+            dict[str -> (float, int)]: mapping from the name of each scalar to the most
+            recent value and the iteration number its added.
+        """
+        return self._latest_scalars
+
+    def latest_with_smoothing_hint(self, window_size=20):
+        """
+        Similar to :meth:`latest`, but the returned values
+        are either the un-smoothed original latest value,
+        or a median of the given window_size,
+        depending on whether the smoothing_hint is True.
+        This provides a default behavior that other writers can use.
+        """
+        result = {}
+        for k, (v, itr) in self._latest_scalars.items():
+            result[k] = (
+                self._history[k].median(window_size) if self._smoothing_hints[k] else v,
+                itr,
+            )
+        return result
+
+    def smoothing_hints(self):
+        """
+        Returns:
+            dict[name -> bool]: the user-provided hint on whether the scalar
+            is noisy and needs smoothing.
+        """
+        return self._smoothing_hints
+
+    def step(self):
+        """
+        User should either: (1) Call this function to increment storage.iter when needed.
+        Or (2) Set `storage.iter` to the correct iteration number before each iteration.
+
+        The storage will then be able to associate the new data with an iteration number.
+        """
+        self._iter += 1
+
+    @property
+    def iter(self):
+        """
+        Returns the current iteration number. When used together with a trainer,
+                this is ensured to be the same as trainer.iter.
+        """
+        return self._iter
+
+    @iter.setter
+    def iter(self, val):
+        self._iter = int(val)
+
+    @property
+    def samples(self):
+        return self._samples
+
+    @samples.setter
+    def samples(self, val):
+        self._samples = int(val)
+
+    def __enter__(self):
+        _CURRENT_STORAGE_STACK.append(self)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        assert _CURRENT_STORAGE_STACK[-1] == self
+        _CURRENT_STORAGE_STACK.pop()
+
+    @contextmanager
+    def name_scope(self, name):
+        """
+        Yields:
+            A context within which all the events added to this storage
+            will be prefixed by the name scope.
+        """
+        old_prefix = self._current_prefix
+        self._current_prefix = name.rstrip("/") + "/"
+        yield
+        self._current_prefix = old_prefix
+
+    def clear_images(self):
+        """
+        Delete all the stored images for visualization. This should be called
+        after images are written to tensorboard.
+        """
+        self._vis_data = []
+
+    def clear_histograms(self):
+        """
+        Delete all the stored histograms for visualization.
+        This should be called after histograms are written to tensorboard.
+        """
+        self._histograms = []

libai/utils/file_io.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import base64
+import concurrent.futures
+import errno
+import logging
+import os
+import shutil
+import tempfile
+import traceback
+from collections import OrderedDict
+from typing import IO, Any, Callable, Dict, Iterable, List, MutableMapping, Optional, Set, Union
+from urllib.parse import urlparse
+
+import portalocker
+
+from libai.utils.download import download
+from libai.utils.non_blocking_io import NonBlockingIOManager
+
+# --------------------------------------------------------
+# References:
+# https://github.com/facebookresearch/iopath/blob/main/iopath/common/file_io.py
+# https://github.com/facebookresearch/detectron2/blob/main/detectron2/utils/file_io.py
+# --------------------------------------------------------
+
+
+__all__ = ["LazyPath", "PathManager", "get_cache_dir", "file_lock"]
+
+
+def get_cache_dir(cache_dir: Optional[str] = None) -> str:
+    """
+    Returns a default directory to cache static files
+    (usually downloaded from Internet), if None is provided.
+
+    Args:
+        cache_dir (None or str): if not None, will be returned as is.
+            If None, returns the default cache directory as:
+
+        1) $LIBAI_CACHE, if set
+        2) otherwise ~/.oneflow/iopath_cache
+    """
+    if cache_dir is None:
+        cache_dir = os.path.expanduser(os.getenv("LIBAI_CACHE", "~/.oneflow/iopath_cache"))
+    try:
+        g_pathmgr.mkdirs(cache_dir)
+        assert os.access(cache_dir, os.W_OK)
+    except (OSError, AssertionError):
+        tmp_dir = os.path.join(tempfile.gettempdir(), "iopath_cache")
+        logger = logging.getLogger(__name__)
+        logger.warning(f"{cache_dir} is not accessible! Using {tmp_dir} instead!")
+        cache_dir = tmp_dir
+    return cache_dir
+
+
+def file_lock(path: str):  # type: ignore
+    """
+    A file lock. Once entered, it is guaranteed that no one else holds the
+    same lock. Others trying to enter the lock will block for 30 minutes and
+    raise an exception.
+    This is useful to make sure workers don't cache files to the same location.
+    Args:
+        path (str): a path to be locked. This function will create a lock named
+            `path + ".lock"`
+    Examples:
+        filename = "/path/to/file"
+        with file_lock(filename):
+            if not os.path.isfile(filename):
+                do_create_file()
+    """
+    dirname = os.path.dirname(path)
+    try:
+        os.makedirs(dirname, exist_ok=True)
+    except OSError:
+        # makedir is not atomic. Exceptions can happen when multiple workers try
+        # to create the same dir, despite exist_ok=True.
+        # When this happens, we assume the dir is created and proceed to creating
+        # the lock. If failed to create the directory, the next line will raise
+        # exceptions.
+        pass
+    return portalocker.Lock(path + ".lock", timeout=3600)  # type: ignore
+
+
+class LazyPath(os.PathLike):
+    """
+    A path that's lazily evaluated when it's used.
+
+    Users should be careful to not use it like a str, because
+    it behaves differently from a str.
+    Path manipulation functions in Python such as `os.path.*` all accept
+    PathLike objects already.
+
+    It can be materialized to a str using `os.fspath`.
+    """
+
+    def __init__(self, func: Callable[[], str]) -> None:
+        """
+        Args:
+            func: a function that takes no arguments and returns the
+                actual path as a str. It will be called at most once.
+        """
+        self._func = func
+        self._value: Optional[str] = None
+
+    def _get_value(self) -> str:
+        if self._value is None:
+            self._value = self._func()
+        return self._value  # pyre-ignore
+
+    def __fspath__(self) -> str:
+        return self._get_value()
+
+    # before more like a str after evaluated
+    def __getattr__(self, name: str):  # type: ignore
+        if self._value is None:
+            raise AttributeError(f"Uninitialized LazyPath has no attribute: {name}.")
+        return getattr(self._value, name)
+
+    def __getitem__(self, key):  # type: ignore
+        if self._value is None:
+            raise TypeError("Uninitialized LazyPath is not subscriptable.")
+        return self._value[key]  # type: ignore
+
+    def __str__(self) -> str:
+        if self._value is not None:
+            return self._value  # type: ignore
+        else:
+            return super().__str__()
+
+
+class PathHandler:
+    """
+    PathHandler is a base class that defines common I/O functionality for a URI
+    protocol. It routes I/O for a generic URI which may look like "protocol://*"
+    or a canonical filepath "/foo/bar/baz".
+    """
+
+    _strict_kwargs_check = True
+
+    def __init__(
+        self,
+        async_executor: Optional[concurrent.futures.Executor] = None,
+    ) -> None:
+        """
+        When registering a `PathHandler`, the user can optionally pass in a
+        `Executor` to run the asynchronous file operations.
+        NOTE: For regular non-async operations of `PathManager`, there is
+        no need to pass `async_executor`.
+
+        Args:
+            async_executor (optional `Executor`): Used for async file operations.
+                Usage:
+                ```
+                    path_handler = NativePathHandler(async_executor=exe)
+                    path_manager.register_handler(path_handler)
+                ```
+        """
+        self._non_blocking_io_manager = None
+        self._non_blocking_io_executor = async_executor
+
+    def _check_kwargs(self, kwargs: Dict[str, Any]) -> None:
+        """
+        Checks if the given arguments are empty. Throws a ValueError if strict
+        kwargs checking is enabled and args are non-empty. If strict kwargs
+        checking is disabled, only a warning is logged.
+
+        Args:
+            kwargs (Dict[str, Any])
+        """
+        if self._strict_kwargs_check:
+            if len(kwargs) > 0:
+                raise ValueError("Unused arguments: {}".format(kwargs))
+        else:
+            logger = logging.getLogger(__name__)
+            for k, v in kwargs.items():
+                logger.warning("[PathManager] {}={} argument ignored".format(k, v))
+
+    def _get_supported_prefixes(self) -> List[str]:
+        """
+        Returns:
+            List[str]: the list of URI prefixes this PathHandler can support
+        """
+        raise NotImplementedError()
+
+    def _get_local_path(self, path: str, force: bool = False, **kwargs: Any) -> str:
+        """
+        Get a filepath which is compatible with native Python I/O such as `open`
+        and `os.path`.
+        If URI points to a remote resource, this function may download and cache
+        the resource to local disk. In this case, the cache stays on filesystem
+        (under `file_io.get_cache_dir()`) and will be used by a different run.
+        Therefore this function is meant to be used with read-only resources.
+        Args:
+            path (str): A URI supported by this PathHandler
+            force(bool): Forces a download from backend if set to True.
+        Returns:
+            local_path (str): a file path which exists on the local file system
+        """
+        raise NotImplementedError()
+
+    def _copy_from_local(
+        self, local_path: str, dst_path: str, overwrite: bool = False, **kwargs: Any
+    ) -> None:
+        """
+        Copies a local file to the specified URI.
+        If the URI is another local path, this should be functionally identical
+        to copy.
+        Args:
+            local_path (str): a file path which exists on the local file system
+            dst_path (str): A URI supported by this PathHandler
+            overwrite (bool): Bool flag for forcing overwrite of existing URI
+        Returns:
+            status (bool): True on success
+        """
+        raise NotImplementedError()
+
+    def _opent(
+        self, path: str, mode: str = "r", buffering: int = 32, **kwargs: Any
+    ) -> Iterable[Any]:
+        raise NotImplementedError()
+
+    def _open(
+        self, path: str, mode: str = "r", buffering: int = -1, **kwargs: Any
+    ) -> Union[IO[str], IO[bytes]]:
+        """
+        Open a stream to a URI, similar to the built-in `open`.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+            mode (str): Specifies the mode in which the file is opened. It defaults
+                to 'r'.
+            buffering (int): An optional integer used to set the buffering policy.
+                Pass 0 to switch buffering off and an integer >= 1 to indicate the
+                size in bytes of a fixed-size chunk buffer. When no buffering
+                argument is given, the default buffering policy depends on the
+                underlying I/O implementation.
+
+        Returns:
+            file: a file-like object.
+        """
+        raise NotImplementedError()
+
+    def _opena(
+        self,
+        path: str,
+        mode: str = "r",
+        callback_after_file_close: Optional[Callable[[None], None]] = None,
+        buffering: int = -1,
+        **kwargs: Any,
+    ) -> Union[IO[str], IO[bytes]]:
+        """
+        Open a stream to a URI with asynchronous permissions.
+
+        NOTE: Writes to the same path are serialized so they are written in
+        the same order as they were called but writes to distinct paths can
+        happen concurrently.
+
+        Usage (default / without callback function):
+            for n in range(50):
+                results = run_a_large_task(n)
+                with path_manager.opena(uri, "w") as f:
+                    f.write(results)            # Runs in separate thread
+                # Main process returns immediately and continues to next iteration
+            path_manager.async_close()
+
+        Usage (advanced / with callback function):
+            # To write local and then copy to Manifold:
+            def cb():
+                path_manager.copy_from_local(
+                    "checkpoint.pt", "manifold://path/to/bucket"
+                )
+            f = pm.opena("checkpoint.pt", "wb", callback_after_file_close=cb)
+            flow.save({...}, f)
+            f.close()
+
+        Args:
+            ...same args as `_open`...
+            callback_after_file_close (Callable): An optional argument that can
+                be passed to perform operations that depend on the asynchronous
+                writes being completed. The file is first written to the local
+                disk and then the callback is executed.
+            buffering (int): An optional argument to set the buffer size for
+                buffered asynchronous writing.
+
+        Returns:
+            file: a file-like object with asynchronous methods.
+        """
+        # Restrict mode until `NonBlockingIO` has async read feature.
+        valid_modes = {"w", "a", "b"}
+        if not all(m in valid_modes for m in mode):
+            raise ValueError("`opena` mode must be write or append")
+
+        # TODO: Each `PathHandler` should set its own `self._buffered`
+        # parameter and pass that in here. Until then, we assume no
+        # buffering for any storage backend.
+        if not self._non_blocking_io_manager:
+            self._non_blocking_io_manager = NonBlockingIOManager(
+                buffered=False,
+                executor=self._non_blocking_io_executor,
+            )
+
+        try:
+            return self._non_blocking_io_manager.get_non_blocking_io(
+                path=self._get_path_with_cwd(path),
+                io_obj=self._open(path, mode, **kwargs),
+                callback_after_file_close=callback_after_file_close,
+                buffering=buffering,
+            )
+        except ValueError:
+            # When `_strict_kwargs_check = True`, then `open_callable`
+            # will throw a `ValueError`. This generic `_opena` function
+            # does not check the kwargs since it may include any `_open`
+            # args like `encoding`, `ttl`, `has_user_data`, etc.
+            logger = logging.getLogger(__name__)
+            logger.exception(
+                "An exception occurred in `NonBlockingIOManager`. This "
+                "is most likely due to invalid `opena` args. Make sure "
+                "they match the `open` args for the `PathHandler`."
+            )
+            self._async_close()
+
+    def _async_join(self, path: Optional[str] = None, **kwargs: Any) -> bool:
+        """
+        Ensures that desired async write threads are properly joined.
+
+        Args:
+            path (str): Pass in a file path to wait until all asynchronous
+                activity for that path is complete. If no path is passed in,
+                then this will wait until all asynchronous jobs are complete.
+
+        Returns:
+            status (bool): True on success
+        """
+        if not self._non_blocking_io_manager:
+            logger = logging.getLogger(__name__)
+            logger.warning(
+                "This is an async feature. No threads to join because " "`opena` was not used."
+            )
+        self._check_kwargs(kwargs)
+        return self._non_blocking_io_manager._join(self._get_path_with_cwd(path) if path else None)
+
+    def _async_close(self, **kwargs: Any) -> bool:
+        """
+        Closes the thread pool used for the asynchronous operations.
+
+        Returns:
+            status (bool): True on success
+        """
+        if not self._non_blocking_io_manager:
+            logger = logging.getLogger(__name__)
+            logger.warning(
+                "This is an async feature. No threadpool to close because " "`opena` was not used."
+            )
+        self._check_kwargs(kwargs)
+        return self._non_blocking_io_manager._close_thread_pool()
+
+    def _copy(self, src_path: str, dst_path: str, overwrite: bool = False, **kwargs: Any) -> bool:
+        """
+        Copies a source path to a destination path.
+
+        Args:
+            src_path (str): A URI supported by this PathHandler
+            dst_path (str): A URI supported by this PathHandler
+            overwrite (bool): Bool flag for forcing overwrite of existing file
+
+        Returns:
+            status (bool): True on success
+        """
+        raise NotImplementedError()
+
+    def _mv(self, src_path: str, dst_path: str, **kwargs: Any) -> bool:
+        """
+        Moves (renames) a source path to a destination path.
+
+        Args:
+            src_path (str): A URI supported by this PathHandler
+            dst_path (str): A URI supported by this PathHandler
+
+        Returns:
+            status (bool): True on success
+        """
+        raise NotImplementedError()
+
+    def _exists(self, path: str, **kwargs: Any) -> bool:
+        """
+        Checks if there is a resource at the given URI.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+
+        Returns:
+            bool: true if the path exists
+        """
+        raise NotImplementedError()
+
+    def _isfile(self, path: str, **kwargs: Any) -> bool:
+        """
+        Checks if the resource at the given URI is a file.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+
+        Returns:
+            bool: true if the path is a file
+        """
+        raise NotImplementedError()
+
+    def _isdir(self, path: str, **kwargs: Any) -> bool:
+        """
+        Checks if the resource at the given URI is a directory.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+
+        Returns:
+            bool: true if the path is a directory
+        """
+        raise NotImplementedError()
+
+    def _ls(self, path: str, **kwargs: Any) -> List[str]:
+        """
+        List the contents of the directory at the provided URI.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+
+        Returns:
+            List[str]: list of contents in given path
+        """
+        raise NotImplementedError()
+
+    def _mkdirs(self, path: str, **kwargs: Any) -> None:
+        """
+        Recursive directory creation function. Like mkdir(), but makes all
+        intermediate-level directories needed to contain the leaf directory.
+        Similar to the native `os.makedirs`.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+        """
+        raise NotImplementedError()
+
+    def _rm(self, path: str, **kwargs: Any) -> None:
+        """
+        Remove the file (not directory) at the provided URI.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+        """
+        raise NotImplementedError()
+
+    def _symlink(self, src_path: str, dst_path: str, **kwargs: Any) -> bool:
+        """
+        Symlink the src_path to the dst_path
+
+        Args:
+            src_path (str): A URI supported by this PathHandler to symlink from
+            dst_path (str): A URI supported by this PathHandler to symlink to
+        """
+        raise NotImplementedError()
+
+    def _set_cwd(self, path: Union[str, None], **kwargs: Any) -> bool:
+        """
+        Set the current working directory. PathHandler classes prepend the cwd
+        to all URI paths that are handled.
+
+        Args:
+            path (str) or None: A URI supported by this PathHandler. Must be a valid
+                absolute path or None to set the cwd to None.
+
+        Returns:
+            bool: true if cwd was set without errors
+        """
+        raise NotImplementedError()
+
+    def _get_path_with_cwd(self, path: str) -> str:
+        """
+        Default implementation. PathHandler classes that provide a `_set_cwd`
+        feature should also override this `_get_path_with_cwd` method.
+
+        Args:
+            path (str): A URI supported by this PathHandler.
+
+        Returns:
+            path (str): Full path with the cwd attached.
+        """
+        return path
+
+
+class NativePathHandler(PathHandler):
+    """
+    Handles paths that can be accessed using Python native system calls. This
+    handler uses `open()` and `os.*` calls on the given path.
+    """
+
+    _cwd = None
+
+    def _get_local_path(self, path: str, force: bool = False, **kwargs: Any) -> str:
+        self._check_kwargs(kwargs)
+        return os.fspath(path)
+
+    def _copy_from_local(
+        self, local_path: str, dst_path: str, overwrite: bool = False, **kwargs: Any
+    ) -> None:
+        self._check_kwargs(kwargs)
+        local_path = self._get_path_with_cwd(local_path)
+        dst_path = self._get_path_with_cwd(dst_path)
+        assert self._copy(src_path=local_path, dst_path=dst_path, overwrite=overwrite, **kwargs)
+
+    def _open(
+        self,
+        path: str,
+        mode: str = "r",
+        buffering: int = -1,
+        encoding: Optional[str] = None,
+        errors: Optional[str] = None,
+        newline: Optional[str] = None,
+        closefd: bool = True,
+        opener: Optional[Callable] = None,
+        **kwargs: Any,
+    ) -> Union[IO[str], IO[bytes]]:
+        """
+        Open a path.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+            mode (str): Specifies the mode in which the file is opened. It defaults
+                to 'r'.
+            buffering (int): An optional integer used to set the buffering policy.
+                Pass 0 to switch buffering off and an integer >= 1 to indicate the
+                size in bytes of a fixed-size chunk buffer. When no buffering
+                argument is given, the default buffering policy works as follows:
+                    * Binary files are buffered in fixed-size chunks; the size of
+                    the buffer is chosen using a heuristic trying to determine the
+                    underlying device’s “block size” and falling back on
+                    io.DEFAULT_BUFFER_SIZE. On many systems, the buffer will
+                    typically be 4096 or 8192 bytes long.
+            encoding (Optional[str]): the name of the encoding used to decode or
+                encode the file. This should only be used in text mode.
+            errors (Optional[str]): an optional string that specifies how encoding
+                and decoding errors are to be handled. This cannot be used in binary
+                mode.
+            newline (Optional[str]): controls how universal newlines mode works
+                (it only applies to text mode). It can be None, '', '\n', '\r',
+                and '\r\n'.
+            closefd (bool): If closefd is False and a file descriptor rather than
+                a filename was given, the underlying file descriptor will be kept
+                open when the file is closed. If a filename is given closefd must
+                be True (the default) otherwise an error will be raised.
+            opener (Optional[Callable]): A custom opener can be used by passing
+                a callable as opener. The underlying file descriptor for the file
+                object is then obtained by calling opener with (file, flags).
+                opener must return an open file descriptor (passing os.open as opener
+                results in functionality similar to passing None).
+            See https://docs.python.org/3/library/functions.html#open for details.
+
+        Returns:
+            file: a file-like object.
+        """
+        self._check_kwargs(kwargs)
+        return open(  # type: ignore
+            self._get_path_with_cwd(path),
+            mode,
+            buffering=buffering,
+            encoding=encoding,
+            errors=errors,
+            newline=newline,
+            closefd=closefd,
+            opener=opener,
+        )
+
+    def _copy(self, src_path: str, dst_path: str, overwrite: bool = False, **kwargs: Any) -> bool:
+        """
+        Copies a source path to a destination path.
+
+        Args:
+            src_path (str): A URI supported by this PathHandler
+            dst_path (str): A URI supported by this PathHandler
+            overwrite (bool): Bool flag for forcing overwrite of existing file
+
+        Returns:
+            status (bool): True on success
+        """
+        self._check_kwargs(kwargs)
+        src_path = self._get_path_with_cwd(src_path)
+        dst_path = self._get_path_with_cwd(dst_path)
+        if os.path.exists(dst_path) and not overwrite:
+            logger = logging.getLogger(__name__)
+            logger.error("Destination file {} already exists.".format(dst_path))
+            return False
+
+        try:
+            shutil.copyfile(src_path, dst_path)
+            return True
+        except Exception as e:
+            logger = logging.getLogger(__name__)
+            logger.error("Error in file copy - {}".format(str(e)))
+            return False
+
+    def _mv(self, src_path: str, dst_path: str, **kwargs: Any) -> bool:
+        """
+        Moves (renames) a source path to a destination path.
+
+        Args:
+            src_path (str): A URI supported by this PathHandler
+            dst_path (str): A URI supported by this PathHandler
+
+        Returns:
+            status (bool): True on success
+        """
+        self._check_kwargs(kwargs)
+        src_path = self._get_path_with_cwd(src_path)
+        dst_path = self._get_path_with_cwd(dst_path)
+        if os.path.exists(dst_path):
+            logger = logging.getLogger(__name__)
+            logger.error("Destination file {} already exists.".format(dst_path))
+            return False
+
+        try:
+            shutil.move(src_path, dst_path)
+            return True
+        except Exception as e:
+            logger = logging.getLogger(__name__)
+            logger.error("Error in move operation - {}".format(str(e)))
+            return False
+
+    def _symlink(self, src_path: str, dst_path: str, **kwargs: Any) -> bool:
+        """
+        Creates a symlink to the src_path at the dst_path
+
+        Args:
+            src_path (str): A URI supported by this PathHandler
+            dst_path (str): A URI supported by this PathHandler
+
+        Returns:
+            status (bool): True on success
+        """
+        self._check_kwargs(kwargs)
+        src_path = self._get_path_with_cwd(src_path)
+        dst_path = self._get_path_with_cwd(dst_path)
+        logger = logging.getLogger(__name__)
+        if not os.path.exists(src_path):
+            logger.error("Source path {} does not exist".format(src_path))
+            return False
+        if os.path.exists(dst_path):
+            logger.error("Destination path {} already exists.".format(dst_path))
+            return False
+        try:
+            os.symlink(src_path, dst_path)
+            return True
+        except Exception as e:
+            logger.error("Error in symlink - {}".format(str(e)))
+            return False
+
+    def _exists(self, path: str, **kwargs: Any) -> bool:
+        self._check_kwargs(kwargs)
+        return os.path.exists(self._get_path_with_cwd(path))
+
+    def _isfile(self, path: str, **kwargs: Any) -> bool:
+        self._check_kwargs(kwargs)
+        return os.path.isfile(self._get_path_with_cwd(path))
+
+    def _isdir(self, path: str, **kwargs: Any) -> bool:
+        self._check_kwargs(kwargs)
+        return os.path.isdir(self._get_path_with_cwd(path))
+
+    def _ls(self, path: str, **kwargs: Any) -> List[str]:
+        self._check_kwargs(kwargs)
+        return os.listdir(self._get_path_with_cwd(path))
+
+    def _mkdirs(self, path: str, **kwargs: Any) -> None:
+        self._check_kwargs(kwargs)
+        try:
+            os.makedirs(path, exist_ok=True)
+        except OSError as e:
+            # EEXIST it can still happen if multiple processes are creating the dir
+            if e.errno != errno.EEXIST:
+                raise
+
+    def _rm(self, path: str, **kwargs: Any) -> None:
+        self._check_kwargs(kwargs)
+        os.remove(path)
+
+    def _set_cwd(self, path: Union[str, None], **kwargs: Any) -> bool:
+        self._check_kwargs(kwargs)
+        # Remove cwd path if None
+        if path is None:
+            self._cwd = None
+            return True
+
+        # Make sure path is a valid Unix path
+        if not os.path.exists(path):
+            raise ValueError(f"{path} is not a valid Unix path")
+        # Make sure path is an absolute path
+        if not os.path.isabs(path):
+            raise ValueError(f"{path} is not an absolute path")
+        self._cwd = path
+        return True
+
+    def _get_path_with_cwd(self, path: str) -> str:
+        return os.path.normpath(path if not self._cwd else os.path.join(self._cwd, path))
+
+
+class HTTPURLHandler(PathHandler):
+    """
+    Download URLs and cache them to disk.
+    """
+
+    def __init__(self) -> None:
+        self.cache_map: Dict[str, str] = {}
+
+    def _get_supported_prefixes(self) -> List[str]:
+        return ["http://", "https://", "ftp://"]
+
+    def _get_local_path(self, path: str, force: bool = False, **kwargs: Any) -> str:
+        """
+        This implementation downloads the remote resource and caches it locally.
+        The resource will only be downloaded if not previously requested.
+        """
+        self._check_kwargs(kwargs)
+        if force or path not in self.cache_map or not os.path.exists(self.cache_map[path]):
+            logger = logging.getLogger(__name__)
+            parsed_url = urlparse(path)
+            dirname = os.path.join(get_cache_dir(), os.path.dirname(parsed_url.path.lstrip("/")))
+            filename = path.split("/")[-1]
+            cached = os.path.join(dirname, filename)
+            with file_lock(cached):
+                if not os.path.isfile(cached):
+                    logger.info("Downloading {} ...".format(path))
+                    cached = download(path, dirname, filename=filename)
+            logger.info("URL {} cached in {}".format(path, cached))
+            self.cache_map[path] = cached
+        return self.cache_map[path]
+
+    def _open(
+        self, path: str, mode: str = "r", buffering: int = -1, **kwargs: Any
+    ) -> Union[IO[str], IO[bytes]]:
+        """
+        Open a remote HTTP path. The resource is first downloaded and cached
+        locally.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+            mode (str): Specifies the mode in which the file is opened. It defaults
+                to 'r'.
+            buffering (int): Not used for this PathHandler.
+
+        Returns:
+            file: a file-like object.
+        """
+        self._check_kwargs(kwargs)
+        assert mode in ("r", "rb"), "{} does not support open with {} mode".format(
+            self.__class__.__name__, mode
+        )
+        assert (
+            buffering == -1
+        ), f"{self.__class__.__name__} does not support the `buffering` argument"
+        local_path = self._get_local_path(path, force=False)
+        return open(local_path, mode)
+
+
+class OneDrivePathHandler(HTTPURLHandler):
+    """
+    Map OneDrive (short) URLs to direct download links
+    """
+
+    ONE_DRIVE_PREFIX = "https://1drv.ms/u/s!"
+
+    def create_one_drive_direct_download(self, one_drive_url: str) -> str:
+        """
+        Converts a short OneDrive URI into a download link that can be used with wget
+
+        Args:
+            one_drive_url (str): A OneDrive URI supported by this PathHandler
+
+        Returns:
+            result_url (str): A direct download URI for the file
+        """
+        data_b64 = base64.b64encode(bytes(one_drive_url, "utf-8"))
+        data_b64_string = data_b64.decode("utf-8").replace("/", "_").replace("+", "-").rstrip("=")
+        result_url = f"https://api.onedrive.com/v1.0/shares/u!{data_b64_string}/root/content"
+        return result_url
+
+    def _get_supported_prefixes(self) -> List[str]:
+        return [self.ONE_DRIVE_PREFIX]
+
+    def _get_local_path(self, path: str, force: bool = False, **kwargs: Any) -> str:
+        """
+        This implementation downloads the remote resource and caches it locally.
+        The resource will only be downloaded if not previously requested.
+        """
+        logger = logging.getLogger(__name__)
+        direct_url = self.create_one_drive_direct_download(path)
+
+        logger.info(f"URL {path} mapped to direct download link {direct_url}")
+
+        return super()._get_local_path(os.fspath(direct_url), force=force, **kwargs)
+
+
+class PathManagerBase:
+    """
+    A class for users to open generic paths or translate generic paths to file names.
+    path_manager.method(path) will do the following:
+
+    1. Find a handler by checking the prefixes in `self._path_handlers`.
+    2. Call handler.method(path) on the handler that's found
+    """
+
+    def __init__(self) -> None:
+        self._path_handlers: MutableMapping[str, PathHandler] = OrderedDict()
+        """
+        Dict for path prefix to handler
+        """
+
+        self._native_path_handler: PathHandler = NativePathHandler()
+        """
+        A NativePathHandler that works on posix paths. This is used as the fallback.
+        """
+
+        self._cwd: Optional[str] = None
+        """
+        Keeps track of the single cwd (if set).
+        NOTE: Only one PathHandler can have a cwd set at a time.
+        """
+
+        self._async_handlers: Set[PathHandler] = set()
+        """
+        Keeps track of the PathHandler subclasses where `opena` was used so
+        all of the threads can be properly joined when calling
+        `PathManager.join`.
+        """
+
+    def __get_path_handler(self, path: Union[str, os.PathLike]) -> PathHandler:
+        """
+        Finds a PathHandler that supports the given path. Falls back to the native
+        PathHandler if no other handler is found.
+
+        Args:
+            path (str or os.PathLike): URI path to resource
+
+        Returns:
+            handler (PathHandler)
+        """
+        path = os.fspath(path)  # pyre-ignore
+        for p in self._path_handlers.keys():
+            if path.startswith(p):
+                return self._path_handlers[p]
+        return self._native_path_handler
+
+    def opent(
+        self, path: str, mode: str = "r", buffering: int = 32, **kwargs: Any
+    ) -> Iterable[Any]:
+        """
+        Open a tabular data source. Only reading is supported.
+        The opent() returns a Python iterable collection object, compared to
+        bytes/text data with open()
+
+        Args:
+            path (str): A URI supported by this PathHandler
+            mode (str): Specifies the mode in which the file is opened. It defaults
+                to 'r'
+            buffering (int): number of rows fetched and cached
+
+        Returns:
+            An iterable collection object.
+        """
+        return self.__get_path_handler(path)._opent(path, mode, buffering, **kwargs)
+
+    def open(
+        self, path: str, mode: str = "r", buffering: int = -1, **kwargs: Any
+    ) -> Union[IO[str], IO[bytes]]:
+        """
+        Open a stream to a URI, similar to the built-in `open`.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+            mode (str): Specifies the mode in which the file is opened. It defaults
+                to 'r'.
+            buffering (int): An optional integer used to set the buffering policy.
+                Pass 0 to switch buffering off and an integer >= 1 to indicate the
+                size in bytes of a fixed-size chunk buffer. When no buffering
+                argument is given, the default buffering policy depends on the
+                underlying I/O implementation.
+
+        Returns:
+            file: a file-like object.
+        """
+        return self.__get_path_handler(path)._open(  # type: ignore
+            path, mode, buffering=buffering, **kwargs
+        )
+
+    # NOTE: This feature is only implemented for `NativePathHandler` and can
+    # currently only be used in write mode.
+    def opena(
+        self,
+        path: str,
+        mode: str = "r",
+        buffering: int = -1,
+        callback_after_file_close: Optional[Callable[[None], None]] = None,
+        **kwargs: Any,
+    ) -> Union[IO[str], IO[bytes]]:
+        """
+        Open a file with asynchronous permissions. `f.write()` calls (and
+        potentially `f.read()` calls in the future) will be dispatched
+        asynchronously such that the main program can continue running.
+
+        NOTE: Writes to the same path are serialized so they are written in
+        the same order as they were called but writes to distinct paths can
+        happen concurrently.
+
+        Usage (default / without callback function):
+            for n in range(50):
+                results = run_a_large_task(n)
+                # `f` is a file-like object with asynchronous methods
+                with path_manager.opena(uri, "w") as f:
+                    f.write(results)            # Runs in separate thread
+                # Main process returns immediately and continues to next iteration
+            path_manager.async_close()
+
+        Usage (advanced / with callback function):
+            # To asynchronously write to Manifold:
+            def cb():
+                path_manager.copy_from_local(
+                    "checkpoint.pt", "manifold://path/to/bucket"
+                )
+            f = pm.opena("checkpoint.pt", "wb", callback_after_file_close=cb)
+            oneflow.save({...}, f)
+            f.close()
+
+        Args:
+            ...
+            callback_after_file_close (Callable): An optional argument that can
+                be passed to perform operations that depend on the asynchronous
+                writes being completed. The file is first written to the local
+                disk and then the callback is executed.
+
+        Returns:
+            file: a file-like object with asynchronous methods.
+        """
+        non_blocking_io = self.__get_path_handler(path)._opena(
+            path,
+            mode,
+            buffering=buffering,
+            callback_after_file_close=callback_after_file_close,
+            **kwargs,
+        )
+        # Keep track of the path handlers where `opena` is used so that all of the
+        # threads can be properly joined on `PathManager.join`.
+        self._async_handlers.add(self.__get_path_handler(path))
+        return non_blocking_io
+
+    def async_join(self, *paths: str, **kwargs: Any) -> bool:
+        """
+        Ensures that desired async write threads are properly joined.
+
+        Usage:
+            Wait for asynchronous methods operating on specific file paths to
+            complete.
+                async_join("path/to/file1.txt")
+                async_join("path/to/file2.txt", "path/to/file3.txt")
+            Wait for all asynchronous methods to complete.
+                async_join()
+
+        Args:
+            *paths (str): Pass in any number of file paths and `async_join` will wait
+                until all asynchronous activity for those paths is complete. If no
+                paths are passed in, then `async_join` will wait until all asynchronous
+                jobs are complete.
+
+        Returns:
+            status (bool): True on success
+        """
+        success = True
+        if not paths:  # Join all.
+            for handler in self._async_handlers:
+                success = handler._async_join(**kwargs) and success
+        else:  # Join specific paths.
+            for path in paths:
+                success = self.__get_path_handler(path)._async_join(path, **kwargs) and success
+        return success
+
+    def async_close(self, **kwargs: Any) -> bool:
+        """
+        `async_close()` must be called at the very end of any script that uses the
+        asynchronous `opena` feature. This calls `async_join()` first and then closes
+        the thread pool used for the asynchronous operations.
+
+        Returns:
+            status (bool): True on success
+        """
+        success = self.async_join(**kwargs)
+        for handler in self._async_handlers:
+            success = handler._async_close(**kwargs) and success
+        self._async_handlers.clear()
+        return success
+
+    def copy(self, src_path: str, dst_path: str, overwrite: bool = False, **kwargs: Any) -> bool:
+        """
+        Copies a source path to a destination path.
+
+        Args:
+            src_path (str): A URI supported by this PathHandler
+            dst_path (str): A URI supported by this PathHandler
+            overwrite (bool): Bool flag for forcing overwrite of existing file
+
+        Returns:
+            status (bool): True on success
+        """
+
+        # Copying across handlers is not supported.
+        assert self.__get_path_handler(src_path) == self.__get_path_handler(  # type: ignore
+            dst_path
+        )
+        return self.__get_path_handler(src_path)._copy(src_path, dst_path, overwrite, **kwargs)
+
+    def mv(self, src_path: str, dst_path: str, **kwargs: Any) -> bool:
+        """
+        Moves (renames) a source path supported by NativePathHandler to
+        a destination path.
+
+        Args:
+            src_path (str): A URI supported by NativePathHandler
+            dst_path (str): A URI supported by NativePathHandler
+
+        Returns:
+            status (bool): True on success
+
+        Exception:
+            Asserts if both the src and dest paths are not supported by
+            NativePathHandler.
+        """
+
+        # Moving across handlers is not supported.
+        assert self.__get_path_handler(src_path) == self.__get_path_handler(  # type: ignore
+            dst_path
+        ), "Src and dest paths must be supported by the same path handler."
+        return self.__get_path_handler(src_path)._mv(src_path, dst_path, **kwargs)
+
+    def get_local_path(self, path: str, force: bool = False, **kwargs: Any) -> str:
+        """
+        Get a filepath which is compatible with native Python I/O such as `open`
+        and `os.path`.
+
+        If URI points to a remote resource, this function may download and cache
+        the resource to local disk.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+            force(bool): Forces a download from backend if set to True.
+        Returns:
+            local_path (str): a file path which exists on the local file system
+        """
+        path = os.fspath(path)
+        return self.__get_path_handler(path)._get_local_path(  # type: ignore
+            path, force=force, **kwargs
+        )
+
+    def copy_from_local(
+        self, local_path: str, dst_path: str, overwrite: bool = False, **kwargs: Any
+    ) -> None:
+        """
+        Copies a local file to the specified URI.
+        If the URI is another local path, this should be functionally identical
+        to copy.
+
+        Args:
+            local_path (str): a file path which exists on the local file system
+            dst_path (str): A URI supported by this PathHandler
+            overwrite (bool): Bool flag for forcing overwrite of existing URI
+
+        Returns:
+            status (bool): True on success
+        """
+        assert os.path.exists(local_path)
+        return self.__get_path_handler(dst_path)._copy_from_local(
+            local_path=local_path, dst_path=dst_path, overwrite=overwrite, **kwargs
+        )
+
+    def exists(self, path: str, **kwargs: Any) -> bool:
+        """
+        Checks if there is a resource at the given URI.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+
+        Returns:
+            bool: true if the path exists
+        """
+        return self.__get_path_handler(path)._exists(path, **kwargs)  # type: ignore
+
+    def isfile(self, path: str, **kwargs: Any) -> bool:
+        """
+        Checks if there the resource at the given URI is a file.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+
+        Returns:
+            bool: true if the path is a file
+        """
+        return self.__get_path_handler(path)._isfile(path, **kwargs)  # type: ignore
+
+    def isdir(self, path: str, **kwargs: Any) -> bool:
+        """
+        Checks if the resource at the given URI is a directory.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+
+        Returns:
+            bool: true if the path is a directory
+        """
+        return self.__get_path_handler(path)._isdir(path, **kwargs)  # type: ignore
+
+    def ls(self, path: str, **kwargs: Any) -> List[str]:
+        """
+        List the contents of the directory at the provided URI.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+
+        Returns:
+            List[str]: list of contents in given path
+        """
+        return self.__get_path_handler(path)._ls(path, **kwargs)
+
+    def mkdirs(self, path: str, **kwargs: Any) -> None:
+        """
+        Recursive directory creation function. Like mkdir(), but makes all
+        intermediate-level directories needed to contain the leaf directory.
+        Similar to the native `os.makedirs`.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+        """
+        return self.__get_path_handler(path)._mkdirs(path, **kwargs)  # type: ignore
+
+    def rm(self, path: str, **kwargs: Any) -> None:
+        """
+        Remove the file (not directory) at the provided URI.
+
+        Args:
+            path (str): A URI supported by this PathHandler
+        """
+        return self.__get_path_handler(path)._rm(path, **kwargs)  # type: ignore
+
+    def symlink(self, src_path: str, dst_path: str, **kwargs: Any) -> bool:
+        """Symlink the src_path to the dst_path
+
+        Args:
+            src_path (str): A URI supported by this PathHandler to symlink from
+            dst_path (str): A URI supported by this PathHandler to symlink to
+        """
+        # Copying across handlers is not supported.
+        assert self.__get_path_handler(src_path) == self.__get_path_handler(  # type: ignore
+            dst_path
+        )
+        return self.__get_path_handler(src_path)._symlink(src_path, dst_path, **kwargs)
+
+    def set_cwd(self, path: Union[str, None], **kwargs: Any) -> bool:
+        """
+        Set the current working directory. PathHandler classes prepend the cwd
+        to all URI paths that are handled.
+
+        Args:
+            path (str) or None: A URI supported by this PathHandler. Must be a valid
+                absolute Unix path or None to set the cwd to None.
+
+        Returns:
+            bool: true if cwd was set without errors
+        """
+        if path is None and self._cwd is None:
+            return True
+        if self.__get_path_handler(path or self._cwd)._set_cwd(path, **kwargs):  # type: ignore
+            self._cwd = path
+            return True
+        return False
+
+    def register_handler(self, handler: PathHandler, allow_override: bool = False) -> None:
+        """
+        Register a path handler associated with `handler._get_supported_prefixes`
+        URI prefixes.
+
+        Args:
+            handler (PathHandler)
+            allow_override (bool): allow overriding existing handler for prefix
+        """
+        logger = logging.getLogger(__name__)
+        assert isinstance(handler, PathHandler), handler
+
+        # Allow override of `NativePathHandler` which is automatically
+        # instantiated by `PathManager`.
+        if isinstance(handler, NativePathHandler):
+            if allow_override:
+                self._native_path_handler = handler
+            else:
+                raise ValueError(
+                    "`NativePathHandler` is registered by default. Use the "
+                    "`allow_override=True` kwarg to override it."
+                )
+            return
+
+        for prefix in handler._get_supported_prefixes():
+            if prefix not in self._path_handlers:
+                self._path_handlers[prefix] = handler
+                continue
+
+            old_handler_type = type(self._path_handlers[prefix])
+            if allow_override:
+                # if using the global PathManager, show the warnings
+                global g_pathmgr
+                if self == g_pathmgr:
+                    logger.warning(
+                        f"[PathManager] Attempting to register prefix '{prefix}' from "
+                        "the following call stack:\n" + "".join(traceback.format_stack(limit=5))
+                        # show the most recent callstack
+                    )
+                    logger.warning(
+                        f"[PathManager] Prefix '{prefix}' is already registered "
+                        f"by {old_handler_type}. We will override the old handler. "
+                        "To avoid such conflicts, create a project-specific PathManager "
+                        "instead."
+                    )
+                self._path_handlers[prefix] = handler
+            else:
+                raise KeyError(
+                    f"[PathManager] Prefix '{prefix}' already registered by {old_handler_type}!"
+                )
+
+        # Sort path handlers in reverse order so longer prefixes take priority,
+        # eg: http://foo/bar before http://foo
+        self._path_handlers = OrderedDict(
+            sorted(self._path_handlers.items(), key=lambda t: t[0], reverse=True)
+        )
+
+    def set_strict_kwargs_checking(self, enable: bool) -> None:
+        """
+        Toggles strict kwargs checking. If enabled, a ValueError is thrown if any
+        unused parameters are passed to a PathHandler function. If disabled, only
+        a warning is given.
+
+        With a centralized file API, there's a tradeoff of convenience and
+        correctness delegating arguments to the proper I/O layers. An underlying
+        `PathHandler` may support custom arguments which should not be statically
+        exposed on the `PathManager` function. For example, a custom `HTTPURLHandler`
+        may want to expose a `cache_timeout` argument for `open()` which specifies
+        how old a locally cached resource can be before it's refetched from the
+        remote server. This argument would not make sense for a `NativePathHandler`.
+        If strict kwargs checking is disabled, `cache_timeout` can be passed to
+        `PathManager.open` which will forward the arguments to the underlying
+        handler. By default, checking is enabled since it is innately unsafe:
+        multiple `PathHandler`s could reuse arguments with different semantic
+        meanings or types.
+
+        Args:
+            enable (bool)
+        """
+        self._native_path_handler._strict_kwargs_check = enable
+        for handler in self._path_handlers.values():
+            handler._strict_kwargs_check = enable
+
+
+class PathManagerFactory:
+    """
+    PathManagerFactory is the class responsible for creating new PathManager
+    instances and removing them when no longer needed.
+    PathManager can be instantiated directly too, but it is recommended that
+    you use PathManagerFactory to create them.
+    """
+
+    GLOBAL_PATH_MANAGER = "global_path_manager"
+    pm_list = {}
+
+    @staticmethod
+    def get(key=GLOBAL_PATH_MANAGER) -> PathManagerBase:
+        """
+        Get the path manager instance associated with a key.
+        A new instance will be created if there is no existing
+        instance associated with the key passed in.
+
+        Args:
+            key (str):
+        """
+        if key not in PathManagerFactory.pm_list:
+            PathManagerFactory.pm_list[key] = PathManagerBase()
+        return PathManagerFactory.pm_list[key]
+
+    @staticmethod
+    def remove(key):
+        """
+        Remove the path manager instance associated with a key.
+
+        Args:
+            key (str):
+        """
+        if key in PathManagerFactory.pm_list:
+            _pm = PathManagerFactory.pm_list.pop(key)  # noqa
+            del _pm
+
+
+"""
+A global instance of PathManager.
+This global instance is provided for backward compatibility, but it is
+recommended that clients use PathManagerFactory
+"""
+g_pathmgr = PathManagerFactory.get()
+
+
+PathManager = PathManagerBase()
+PathManager.register_handler(HTTPURLHandler())
+PathManager.register_handler(OneDrivePathHandler())

libai/utils/file_utils.py

+"""
+Utilities for working with the local dataset cache.
+This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp
+Copyright by the AllenNLP authors.
+"""
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import fnmatch
+import hashlib
+import json
+import logging
+import os
+import shutil
+import sys
+import tempfile
+from functools import wraps
+from io import open
+from pathlib import Path
+
+import boto3
+import requests
+import wget
+from botocore.config import Config
+from botocore.exceptions import ClientError
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+
+cache_home = Path(os.getenv("OF_CACHE_ROOT", Path.home() / ".of_cache"))
+default_cache_path = str(cache_home / "libai")
+
+try:
+    from urllib.parse import urlparse
+except ImportError:
+    from urlparse import urlparse
+
+
+def url_to_filename(url, etag=None):
+    """
+    Convert `url` into a hashed filename in a repeatable way.
+    If `etag` is specified, append its hash to the url's, delimited
+    by a period.
+    If the url ends with .h5 (Keras HDF5 weights) ands '.h5' to the name
+    so that TF 2.0 can identify it as a HDF5 file
+    (see https://github.com/tensorflow/tensorflow/blob/00fad90125b18b80fe054de1055770cfb8fe4ba3
+    /tensorflow/python/keras/engine/network.py#L1380)
+    """
+    url_bytes = url.encode("utf-8")
+    url_hash = hashlib.sha256(url_bytes)
+    filename = url_hash.hexdigest()
+
+    if etag:
+        etag_bytes = etag.encode("utf-8")
+        etag_hash = hashlib.sha256(etag_bytes)
+        filename += "." + etag_hash.hexdigest()
+
+    if url.endswith(".h5"):
+        filename += ".h5"
+
+    return filename
+
+
+def filename_to_url(filename, cache_dir=None):
+    """
+    Return the url and etag (which may be ``None``) stored for `filename`.
+    Raise ``EnvironmentError`` if `filename` or its stored metadata do not exist.
+    """
+    if cache_dir is None:
+        cache_dir = default_cache_path
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    cache_path = os.path.join(cache_dir, filename)
+    if not os.path.exists(cache_path):
+        raise EnvironmentError("file {} not found".format(cache_path))
+
+    meta_path = cache_path + ".json"
+    if not os.path.exists(meta_path):
+        raise EnvironmentError("file {} not found".format(meta_path))
+
+    with open(meta_path, encoding="utf-8") as meta_file:
+        metadata = json.load(meta_file)
+    url = metadata["url"]
+    etag = metadata["etag"]
+
+    return url, etag
+
+
+def cached_path(url_or_filename, cache_dir=None, force_download=False, proxies=None):
+    """
+    Given something that might be a URL (or might be a local path),
+    determine which. If it's a URL, download the file and cache it, and
+    return the path to the cached file. If it's already a local path,
+    make sure the file exists and then return the path.
+    Args:
+        cache_dir: specify a cache directory to save the file to (overwrite the default cache dir).
+        force_download: if True, re-dowload the file even if it's already cached in the cache dir.
+    """
+    if cache_dir is None:
+        cache_dir = default_cache_path
+    if sys.version_info[0] == 3 and isinstance(url_or_filename, Path):
+        url_or_filename = str(url_or_filename)
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    parsed = urlparse(url_or_filename)
+
+    if parsed.scheme in ("http", "https", "s3"):
+        # URL, so get it from the cache (downloading if necessary)
+        return get_from_cache(
+            url_or_filename,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+        )
+    elif os.path.exists(url_or_filename):
+        # File, and it exists.
+        return url_or_filename
+    elif parsed.scheme == "":
+        # File, but it doesn't exist.
+        raise EnvironmentError("file {} not found".format(url_or_filename))
+    else:
+        # Something unknown
+        raise ValueError("unable to parse {} as a URL or as a local path".format(url_or_filename))
+
+
+def split_s3_path(url):
+    """Split a full s3 path into the bucket name and path."""
+    parsed = urlparse(url)
+    if not parsed.netloc or not parsed.path:
+        raise ValueError("bad s3 path {}".format(url))
+    bucket_name = parsed.netloc
+    s3_path = parsed.path
+    # Remove '/' at beginning of path.
+    if s3_path.startswith("/"):
+        s3_path = s3_path[1:]
+    return bucket_name, s3_path
+
+
+def s3_request(func):
+    """
+    Wrapper function for s3 requests in order to create more helpful error
+    messages.
+    """
+
+    @wraps(func)
+    def wrapper(url, *args, **kwargs):
+        try:
+            return func(url, *args, **kwargs)
+        except ClientError as exc:
+            if int(exc.response["Error"]["Code"]) == 404:
+                raise EnvironmentError("file {} not found".format(url))
+            else:
+                raise
+
+    return wrapper
+
+
+@s3_request
+def s3_etag(url, proxies=None):
+    """Check ETag on S3 object."""
+    s3_resource = boto3.resource("s3", config=Config(proxies=proxies))
+    bucket_name, s3_path = split_s3_path(url)
+    s3_object = s3_resource.Object(bucket_name, s3_path)
+    return s3_object.e_tag
+
+
+@s3_request
+def s3_get(url, temp_file, proxies=None):
+    """Pull a file directly from S3."""
+    s3_resource = boto3.resource("s3", config=Config(proxies=proxies))
+    bucket_name, s3_path = split_s3_path(url)
+    s3_resource.Bucket(bucket_name).download_fileobj(s3_path, temp_file)
+
+
+def http_get(url, temp_file, proxies=None):
+    req = requests.get(url, stream=True, proxies=proxies)
+    content_length = req.headers.get("Content-Length")
+    total = int(content_length) if content_length is not None else None
+    progress = tqdm(unit="B", total=total)
+    for chunk in req.iter_content(chunk_size=1024):
+        if chunk:  # filter out keep-alive new chunks
+            progress.update(len(chunk))
+            temp_file.write(chunk)
+    progress.close()
+
+
+def get_from_cache(url, cache_dir=None, force_download=False, proxies=None, etag_timeout=10):
+    """
+    Given a URL, look for the corresponding dataset in the local cache.
+    If it's not there, download it. Then return the path to the cached file.
+    """
+    if cache_dir is None:
+        cache_dir = default_cache_path
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    if not os.path.exists(cache_dir):
+        os.makedirs(cache_dir)
+
+    # Get eTag to add to filename, if it exists.
+    if url.startswith("s3://"):
+        etag = s3_etag(url, proxies=proxies)
+    else:
+        try:
+            response = requests.head(
+                url, allow_redirects=True, proxies=proxies, timeout=etag_timeout
+            )
+            if response.status_code != 200:
+                etag = None
+            else:
+                etag = response.headers.get("ETag")
+        except (EnvironmentError, requests.exceptions.Timeout):
+            etag = None
+
+    filename = url_to_filename(url, etag)
+
+    # get cache path to put the file
+    cache_path = os.path.join(cache_dir, filename)
+
+    # If we don't have a connection (etag is None) and can't identify the file
+    # try to get the last downloaded one
+    if not os.path.exists(cache_path) and etag is None:
+        matching_files = fnmatch.filter(os.listdir(cache_dir), filename + ".*")
+        matching_files = list(filter(lambda s: not s.endswith(".json"), matching_files))
+        if matching_files:
+            cache_path = os.path.join(cache_dir, matching_files[-1])
+
+    if not os.path.exists(cache_path) or force_download:
+        # Download to temporary file, then copy to cache dir once finished.
+        # Otherwise you get corrupt cache entries if the download gets interrupted.
+        with tempfile.NamedTemporaryFile() as temp_file:
+            logger.info(
+                "%s not found in cache or force_download set to True, downloading to %s",
+                url,
+                temp_file.name,
+            )
+
+            # GET file object
+            if url.startswith("s3://"):
+                s3_get(url, temp_file, proxies=proxies)
+            else:
+                http_get(url, temp_file, proxies=proxies)
+
+            # we are copying the file before closing it, so flush to avoid truncation
+            temp_file.flush()
+            # shutil.copyfileobj() starts at the current position, so go to the start
+            temp_file.seek(0)
+
+            logger.info("copying %s to cache at %s", temp_file.name, cache_path)
+            with open(cache_path, "wb") as cache_file:
+                shutil.copyfileobj(temp_file, cache_file)
+
+            logger.info("creating metadata file for %s", cache_path)
+            meta = {"url": url, "etag": etag}
+            meta_path = cache_path + ".json"
+            with open(meta_path, "w") as meta_file:
+                output_string = json.dumps(meta)
+                meta_file.write(output_string)
+
+            logger.info("removing temp file %s", temp_file.name)
+
+    return cache_path
+
+
+def get_md5(fname):
+    hash_md5 = hashlib.md5()
+    with open(fname, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_md5.update(chunk)
+    result = hash_md5.hexdigest()
+    return result
+
+
+def download_file(out_path: str, url):
+    logger.info(f"downloading from {url} to {out_path}")
+    wget.download(url, out=out_path)
+
+
+def get_data_from_cache(url, cache_dir=None, force_download=False, md5=None):
+    """
+    Given a URL, look for the corresponding dataset in the local cache.
+    If it's not there, download it. Then return the path to the cached file.
+    """
+    if cache_dir is None:
+        cache_dir = default_cache_path
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    if not os.path.exists(cache_dir):
+        os.makedirs(cache_dir)
+
+    filename = url.split("/")[-1]
+
+    # get cache path to put the file
+    cache_path = os.path.join(cache_dir, filename)
+
+    # If we have already get the file, just check the md5 if provided
+    if os.path.exists(cache_path) and md5 is not None:
+        local_file_md5 = get_md5(cache_path)
+        if local_file_md5 != md5:
+            os.unlink(cache_path)
+            download_file(cache_path, url)
+
+    # If we don't have a connection (etag is None) and can't identify the file
+    # try to get the last downloaded one
+    if not os.path.exists(cache_path):
+        download_file(cache_path, url)
+
+    if not os.path.exists(cache_path) or force_download:
+        # Download to temporary file, then copy to cache dir once finished.
+        # Otherwise you get corrupt cache entries if the download gets interrupted.
+        with tempfile.NamedTemporaryFile() as temp_file:
+            logger.info(
+                "%s not found in cache or force_download set to True, downloading to %s",
+                url,
+                temp_file.name,
+            )
+
+            # GET file object
+            http_get(url, temp_file)
+
+            # we are copying the file before closing it, so flush to avoid truncation
+            temp_file.flush()
+            # shutil.copyfileobj() starts at the current position, so go to the start
+            temp_file.seek(0)
+
+            logger.info("copying %s to cache at %s", temp_file.name, cache_path)
+            with open(cache_path, "wb") as cache_file:
+                shutil.copyfileobj(temp_file, cache_file)
+
+            logger.info("creating metadata file for %s", cache_path)
+            meta = {"url": url}
+            meta_path = cache_path + ".json"
+            with open(meta_path, "w") as meta_file:
+                output_string = json.dumps(meta)
+                meta_file.write(output_string)
+
+            logger.info("removing temp file %s", temp_file.name)
+
+    return cache_path

libai/utils/history_buffer.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Tuple
+
+import numpy as np
+
+# --------------------------------------------------------
+# References:
+# https://github.com/facebookresearch/fvcore/blob/main/fvcore/common/history_buffer.py
+# --------------------------------------------------------
+
+
+class HistoryBuffer:
+    """
+    Track a series of scalar values and provide access to smoothed values over a
+    window or the global average of the series.
+    """
+
+    def __init__(self, max_length: int = 1000000):
+        """
+        Args:
+            max_length: maximal number of values that can be stored in the
+                buffer. When the capacity of the buffer is exhausted, old
+                values will be removed.
+        """
+        self._max_length: int = max_length
+        self._data: List[Tuple[float, float]] = []  # (value, iteration) pairs
+        self._count: int = 0
+        self._global_avg: float = 0
+
+    def update(self, value: float, iteration: float = None):
+        """
+        Add a new scalar value produced at certain iteration. If the length
+        of the buffer exceeds self._max_length, the oldest element will be
+        removed from the buffer.
+        """
+        if iteration is None:
+            iteration = self._count
+        if len(self._data) == self._max_length:
+            self._data.pop(0)
+        self._data.append((value, iteration))
+
+        self._count += 1
+        self._global_avg += (value - self._global_avg) / self._count
+
+    def latest(self):
+        """
+        Return the latest scalar value added to the buffer.
+        """
+        return self._data[-1][0]
+
+    def median(self, window_size: int):
+        """
+        Return the median of the latest `window_size` values in the buffer.
+        """
+        return np.median([x[0] for x in self._data[-window_size:]])
+
+    def avg(self, window_size: int):
+        """
+        Return the mean of the latest `window_size` values in the buffer.
+        """
+        return np.mean([x[0] for x in self._data[-window_size:]])
+
+    def global_avg(self):
+        """
+        Return the mean of all the elements in the buffer. Note that this
+        includes those getting removed due to limited buffer storage.
+        """
+        return self._global_avg
+
+    def values(self):
+        """
+        Returns:
+            list[(number, iteration)]: content of the current buffer.
+        """
+        return self._data