File size: 1,432 Bytes
b170003 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
# !/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE_TaBERT file in the root directory of this source tree.
import logging
from enum import Enum
class TransformerVersion(Enum):
PYTORCH_PRETRAINED_BERT = 0
TRANSFORMERS = 1
TRANSFORMER_VERSION = None
try:
from pytorch_pretrained_bert.modeling import (
BertForMaskedLM, BertForPreTraining, BertModel,
BertConfig,
BertSelfOutput, BertIntermediate, BertOutput,
BertLMPredictionHead, BertLayerNorm, gelu
)
from pytorch_pretrained_bert.tokenization import BertTokenizer
hf_flag = 'old'
TRANSFORMER_VERSION = TransformerVersion.PYTORCH_PRETRAINED_BERT
logging.warning('You are using the old version of `pytorch_pretrained_bert`')
except Exception as e:
print(f"Using transformers as {e}")
from transformers.models.bert import BertTokenizer # noqa
from transformers import ( # noqa
BertForMaskedLM, BertForPreTraining, BertModel, BertConfig
)
from transformers.models.bert.modeling_bert import (
BertSelfOutput, BertIntermediate, BertOutput,
BertLMPredictionHead
)
from transformers.activations import gelu
from torch.nn import LayerNorm as BertLayerNorm
hf_flag = 'new'
TRANSFORMER_VERSION = TransformerVersion.TRANSFORMERS
|