__init__.py 4.47 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
__version__ = "1.2.0"
thomwolf's avatar
thomwolf committed
2
3
4
5
6
7
8
9
10
11
12
# Work around to update TensorFlow's absl.logging threshold which alters the
# default Python logging output behavior when present.
# see: https://github.com/abseil/abseil-py/issues/99
# and: https://github.com/tensorflow/tensorflow/issues/26691#issuecomment-500369493
try:
    import absl.logging
    absl.logging.set_verbosity('info')
    absl.logging.set_stderrthreshold('info')
    absl.logging._warn_preinit_stderr = False
except:
    pass
13
14
15

# Tokenizer
from .tokenization_utils import (PreTrainedTokenizer)
thomwolf's avatar
thomwolf committed
16
from .tokenization_auto import AutoTokenizer
thomwolf's avatar
thomwolf committed
17
from .tokenization_bert import BertTokenizer, BasicTokenizer, WordpieceTokenizer
thomwolf's avatar
thomwolf committed
18
from .tokenization_openai import OpenAIGPTTokenizer
19
from .tokenization_transfo_xl import (TransfoXLTokenizer, TransfoXLCorpus)
thomwolf's avatar
thomwolf committed
20
from .tokenization_gpt2 import GPT2Tokenizer
21
from .tokenization_xlnet import XLNetTokenizer, SPIECE_UNDERLINE
thomwolf's avatar
thomwolf committed
22
from .tokenization_xlm import XLMTokenizer
23
from .tokenization_roberta import RobertaTokenizer
thomwolf's avatar
thomwolf committed
24
from .tokenization_distilbert import DistilBertTokenizer
25

26
# Configurations
27
from .configuration_utils import PretrainedConfig
28
29
30
31
32
33
34
35
36
from .configuration_auto import AutoConfig
from .configuration_bert import BertConfig, BERT_PRETRAINED_CONFIG_ARCHIVE_MAP
from .configuration_openai import OpenAIGPTConfig, OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP
from .configuration_transfo_xl import TransfoXLConfig, TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP
from .configuration_gpt2 import GPT2Config, GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP
from .configuration_xlnet import XLNetConfig, XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP
from .configuration_xlm import XLMConfig, XLM_PRETRAINED_CONFIG_ARCHIVE_MAP
from .configuration_roberta import RobertaConfig, ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP
from .configuration_distilbert import DistilBertConfig, DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP
37

38
# Modeling
39
from .modeling_utils import (PreTrainedModel, prune_layer, Conv1D)
40
from .modeling_auto import (AutoModel, AutoModelForSequenceClassification, AutoModelForQuestionAnswering,
41
                            AutoModelWithLMHead)
thomwolf's avatar
thomwolf committed
42

43
from .modeling_bert import (BertPreTrainedModel, BertModel, BertForPreTraining,
thomwolf's avatar
thomwolf committed
44
45
46
                            BertForMaskedLM, BertForNextSentencePrediction,
                            BertForSequenceClassification, BertForMultipleChoice,
                            BertForTokenClassification, BertForQuestionAnswering,
47
48
                            load_tf_weights_in_bert, BERT_PRETRAINED_MODEL_ARCHIVE_MAP)
from .modeling_openai import (OpenAIGPTPreTrainedModel, OpenAIGPTModel,
thomwolf's avatar
thomwolf committed
49
                              OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel,
50
51
52
53
                              load_tf_weights_in_openai_gpt, OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_MAP)
from .modeling_transfo_xl import (TransfoXLPreTrainedModel, TransfoXLModel, TransfoXLLMHeadModel,
                                  load_tf_weights_in_transfo_xl, TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_MAP)
from .modeling_gpt2 import (GPT2PreTrainedModel, GPT2Model,
thomwolf's avatar
thomwolf committed
54
                            GPT2LMHeadModel, GPT2DoubleHeadsModel,
55
56
                            load_tf_weights_in_gpt2, GPT2_PRETRAINED_MODEL_ARCHIVE_MAP)
from .modeling_xlnet import (XLNetPreTrainedModel, XLNetModel, XLNetLMHeadModel,
57
                             XLNetForSequenceClassification, XLNetForQuestionAnswering, XLNetForMultipleChoice,
58
59
                             load_tf_weights_in_xlnet, XLNET_PRETRAINED_MODEL_ARCHIVE_MAP)
from .modeling_xlm import (XLMPreTrainedModel , XLMModel,
thomwolf's avatar
thomwolf committed
60
                           XLMWithLMHeadModel, XLMForSequenceClassification,
61
62
                           XLMForQuestionAnswering, XLM_PRETRAINED_MODEL_ARCHIVE_MAP)
from .modeling_roberta import (RobertaForMaskedLM, RobertaModel, RobertaForSequenceClassification,
63
                               RobertaForMultipleChoice, ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP)
64
from .modeling_distilbert import (DistilBertForMaskedLM, DistilBertModel,
thomwolf's avatar
thomwolf committed
65
                               DistilBertForSequenceClassification, DistilBertForQuestionAnswering,
66
                               DISTILBERT_PRETRAINED_MODEL_ARCHIVE_MAP)
67

68
# Optimization
thomwolf's avatar
thomwolf committed
69
70
from .optimization import (AdamW, ConstantLRSchedule, WarmupConstantSchedule, WarmupCosineSchedule,
                           WarmupCosineWithHardRestartsSchedule, WarmupLinearSchedule)
71

72
# Files and general utilities
73
74
75
from .file_utils import (PYTORCH_TRANSFORMERS_CACHE, PYTORCH_PRETRAINED_BERT_CACHE,
                         cached_path, add_start_docstrings, add_end_docstrings,
                         WEIGHTS_NAME, TF_WEIGHTS_NAME, CONFIG_NAME)