__init__.py 30.3 KB
Newer Older
1
2
3
4
# flake8: noqa
# There's no way to ignore "F401 '...' imported but unused" warnings in this
# module, but to preserve other warnings. So, don't check this module at all.

5
__version__ = "4.0.0-dev"
thomwolf's avatar
thomwolf committed
6

thomwolf's avatar
thomwolf committed
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
13
14
15
except ImportError:
    pass
else:
16
17
    absl.logging.set_verbosity("info")
    absl.logging.set_stderrthreshold("info")
thomwolf's avatar
thomwolf committed
18
    absl.logging._warn_preinit_stderr = False
19

20
21
22
23
24
25
26
27
28
29
# Integrations: this needs to come before other ml imports
# in order to allow any 3rd-party code to initialize properly
from .integrations import (  # isort:skip
    is_comet_available,
    is_optuna_available,
    is_ray_available,
    is_tensorboard_available,
    is_wandb_available,
)

Sylvain Gugger's avatar
Sylvain Gugger committed
30
# Configuration
Aymeric Augustin's avatar
Aymeric Augustin committed
31
from .configuration_utils import PretrainedConfig
Sylvain Gugger's avatar
Sylvain Gugger committed
32
33

# Data
34
from .data import (
Aymeric Augustin's avatar
Aymeric Augustin committed
35
    DataProcessor,
36
37
38
    InputExample,
    InputFeatures,
    SingleSentenceClassificationProcessor,
Aymeric Augustin's avatar
Aymeric Augustin committed
39
40
41
42
    SquadExample,
    SquadFeatures,
    SquadV1Processor,
    SquadV2Processor,
43
    glue_compute_metrics,
44
    glue_convert_examples_to_features,
Aymeric Augustin's avatar
Aymeric Augustin committed
45
    glue_output_modes,
46
47
    glue_processors,
    glue_tasks_num_labels,
Aymeric Augustin's avatar
Aymeric Augustin committed
48
    squad_convert_examples_to_features,
49
    xnli_compute_metrics,
50
51
52
53
    xnli_output_modes,
    xnli_processors,
    xnli_tasks_num_labels,
)
54

Aymeric Augustin's avatar
Aymeric Augustin committed
55
56
57
58
59
60
# Files and general utilities
from .file_utils import (
    CONFIG_NAME,
    MODEL_CARD_NAME,
    PYTORCH_PRETRAINED_BERT_CACHE,
    PYTORCH_TRANSFORMERS_CACHE,
61
    SPIECE_UNDERLINE,
Aymeric Augustin's avatar
Aymeric Augustin committed
62
63
64
65
66
67
68
    TF2_WEIGHTS_NAME,
    TF_WEIGHTS_NAME,
    TRANSFORMERS_CACHE,
    WEIGHTS_NAME,
    add_end_docstrings,
    add_start_docstrings,
    cached_path,
Patrick von Platen's avatar
Patrick von Platen committed
69
    is_apex_available,
70
    is_datasets_available,
Ola Piktus's avatar
Ola Piktus committed
71
    is_faiss_available,
72
    is_flax_available,
Patrick von Platen's avatar
Patrick von Platen committed
73
74
    is_psutil_available,
    is_py3nvml_available,
75
    is_sentencepiece_available,
76
    is_sklearn_available,
Aymeric Augustin's avatar
Aymeric Augustin committed
77
    is_tf_available,
78
    is_tokenizers_available,
Aymeric Augustin's avatar
Aymeric Augustin committed
79
    is_torch_available,
80
    is_torch_tpu_available,
Aymeric Augustin's avatar
Aymeric Augustin committed
81
)
82
from .hf_argparser import HfArgumentParser
83

thomwolf's avatar
thomwolf committed
84
# Model Cards
85
from .modelcard import ModelCard
86

Aymeric Augustin's avatar
Aymeric Augustin committed
87
88
89
90
91
92
93
94
95
96
# TF 2.0 <=> PyTorch conversion utilities
from .modeling_tf_pytorch_utils import (
    convert_tf_weight_name_to_pt_weight_name,
    load_pytorch_checkpoint_in_tf2_model,
    load_pytorch_model_in_tf2_model,
    load_pytorch_weights_in_tf2_model,
    load_tf2_checkpoint_in_pytorch_model,
    load_tf2_model_in_pytorch_model,
    load_tf2_weights_in_pytorch_model,
)
Sylvain Gugger's avatar
Sylvain Gugger committed
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
from .models.albert import ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, AlbertConfig
from .models.auto import (
    ALL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    CONFIG_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoTokenizer,
)
from .models.bart import BartConfig, BartTokenizer
from .models.bert import (
    BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    BasicTokenizer,
    BertConfig,
    BertTokenizer,
    WordpieceTokenizer,
)
from .models.bert_generation import BertGenerationConfig
from .models.bert_japanese import BertJapaneseTokenizer, CharacterTokenizer, MecabTokenizer
from .models.bertweet import BertweetTokenizer
from .models.blenderbot import (
    BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    BlenderbotConfig,
    BlenderbotSmallTokenizer,
    BlenderbotTokenizer,
)
from .models.camembert import CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, CamembertConfig
from .models.ctrl import CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP, CTRLConfig, CTRLTokenizer
from .models.deberta import DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, DebertaConfig, DebertaTokenizer
from .models.distilbert import DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, DistilBertConfig, DistilBertTokenizer
from .models.dpr import (
    DPR_PRETRAINED_CONFIG_ARCHIVE_MAP,
    DPRConfig,
    DPRContextEncoderTokenizer,
    DPRQuestionEncoderTokenizer,
    DPRReaderOutput,
    DPRReaderTokenizer,
)
from .models.electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraTokenizer
from .models.encoder_decoder import EncoderDecoderConfig
from .models.flaubert import FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, FlaubertConfig, FlaubertTokenizer
from .models.fsmt import FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP, FSMTConfig, FSMTTokenizer
from .models.funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig, FunnelTokenizer
from .models.gpt2 import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPT2Config, GPT2Tokenizer
from .models.herbert import HerbertTokenizer
from .models.layoutlm import LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMConfig, LayoutLMTokenizer
from .models.longformer import LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, LongformerConfig, LongformerTokenizer
from .models.lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig, LxmertTokenizer
from .models.marian import MarianConfig
from .models.mbart import MBartConfig
from .models.mmbt import MMBTConfig
from .models.mobilebert import MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MobileBertConfig, MobileBertTokenizer
Sylvain Gugger's avatar
Sylvain Gugger committed
148
from .models.mt5 import MT5Config
Sylvain Gugger's avatar
Sylvain Gugger committed
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
from .models.openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig, OpenAIGPTTokenizer
from .models.pegasus import PegasusConfig
from .models.phobert import PhobertTokenizer
from .models.prophetnet import PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP, ProphetNetConfig, ProphetNetTokenizer
from .models.rag import RagConfig, RagRetriever, RagTokenizer
from .models.reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig
from .models.retribert import RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RetriBertConfig, RetriBertTokenizer
from .models.roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig, RobertaTokenizer
from .models.squeezebert import SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, SqueezeBertConfig, SqueezeBertTokenizer
from .models.t5 import T5_PRETRAINED_CONFIG_ARCHIVE_MAP, T5Config
from .models.transfo_xl import (
    TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    TransfoXLConfig,
    TransfoXLCorpus,
    TransfoXLTokenizer,
)
from .models.xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMTokenizer
from .models.xlm_prophetnet import XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMProphetNetConfig
from .models.xlm_roberta import XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMRobertaConfig
from .models.xlnet import XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLNetConfig
169

Aymeric Augustin's avatar
Aymeric Augustin committed
170
171
# Pipelines
from .pipelines import (
172
173
    Conversation,
    ConversationalPipeline,
Aymeric Augustin's avatar
Aymeric Augustin committed
174
175
    CsvPipelineDataFormat,
    FeatureExtractionPipeline,
Julien Chaumond's avatar
Julien Chaumond committed
176
    FillMaskPipeline,
Aymeric Augustin's avatar
Aymeric Augustin committed
177
178
179
180
181
182
    JsonPipelineDataFormat,
    NerPipeline,
    PipedPipelineDataFormat,
    Pipeline,
    PipelineDataFormat,
    QuestionAnsweringPipeline,
183
    SummarizationPipeline,
184
    Text2TextGenerationPipeline,
Aymeric Augustin's avatar
Aymeric Augustin committed
185
    TextClassificationPipeline,
186
    TextGenerationPipeline,
187
    TokenClassificationPipeline,
188
    TranslationPipeline,
189
    ZeroShotClassificationPipeline,
Aymeric Augustin's avatar
Aymeric Augustin committed
190
191
    pipeline,
)
Julien Chaumond's avatar
Julien Chaumond committed
192

Sylvain Gugger's avatar
Sylvain Gugger committed
193
# Tokenization
194
195
from .tokenization_utils import PreTrainedTokenizer
from .tokenization_utils_base import (
196
    AddedToken,
197
    BatchEncoding,
198
199
    CharSpan,
    PreTrainedTokenizerBase,
200
201
    SpecialTokensMixin,
    TensorType,
202
    TokenSpan,
203
)
204
205
206


if is_sentencepiece_available():
Sylvain Gugger's avatar
Sylvain Gugger committed
207
208
209
210
211
212
213
214
215
216
217
    from .models.albert import AlbertTokenizer
    from .models.bert_generation import BertGenerationTokenizer
    from .models.camembert import CamembertTokenizer
    from .models.marian import MarianTokenizer
    from .models.mbart import MBartTokenizer
    from .models.pegasus import PegasusTokenizer
    from .models.reformer import ReformerTokenizer
    from .models.t5 import T5Tokenizer
    from .models.xlm_prophetnet import XLMProphetNetTokenizer
    from .models.xlm_roberta import XLMRobertaTokenizer
    from .models.xlnet import XLNetTokenizer
218
219
220
221
else:
    from .utils.dummy_sentencepiece_objects import *

if is_tokenizers_available():
Sylvain Gugger's avatar
Sylvain Gugger committed
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
    from .models.albert import AlbertTokenizerFast
    from .models.bart import BartTokenizerFast
    from .models.bert import BertTokenizerFast
    from .models.camembert import CamembertTokenizerFast
    from .models.distilbert import DistilBertTokenizerFast
    from .models.dpr import DPRContextEncoderTokenizerFast, DPRQuestionEncoderTokenizerFast, DPRReaderTokenizerFast
    from .models.electra import ElectraTokenizerFast
    from .models.funnel import FunnelTokenizerFast
    from .models.gpt2 import GPT2TokenizerFast
    from .models.herbert import HerbertTokenizerFast
    from .models.layoutlm import LayoutLMTokenizerFast
    from .models.longformer import LongformerTokenizerFast
    from .models.lxmert import LxmertTokenizerFast
    from .models.mbart import MBartTokenizerFast
    from .models.mobilebert import MobileBertTokenizerFast
    from .models.openai import OpenAIGPTTokenizerFast
    from .models.pegasus import PegasusTokenizerFast
    from .models.reformer import ReformerTokenizerFast
    from .models.retribert import RetriBertTokenizerFast
    from .models.roberta import RobertaTokenizerFast
    from .models.squeezebert import SqueezeBertTokenizerFast
    from .models.t5 import T5TokenizerFast
    from .models.xlm_roberta import XLMRobertaTokenizerFast
    from .models.xlnet import XLNetTokenizerFast
246
247
248
249
250
251
252
    from .tokenization_utils_fast import PreTrainedTokenizerFast

    if is_sentencepiece_available():
        from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS, convert_slow_tokenizer
else:
    from .utils.dummy_tokenizers_objects import *

253
# Trainer
Sylvain Gugger's avatar
Sylvain Gugger committed
254
255
256
257
258
259
260
261
262
from .trainer_callback import (
    DefaultFlowCallback,
    PrinterCallback,
    ProgressCallback,
    TrainerCallback,
    TrainerControl,
    TrainerState,
)
from .trainer_utils import EvalPrediction, EvaluationStrategy, set_seed
263
from .training_args import TrainingArguments
Julien Plu's avatar
Julien Plu committed
264
from .training_args_tf import TFTrainingArguments
Lysandre Debut's avatar
Lysandre Debut committed
265
from .utils import logging
Aymeric Augustin's avatar
Aymeric Augustin committed
266
267


Lysandre Debut's avatar
Lysandre Debut committed
268
logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
Aymeric Augustin's avatar
Aymeric Augustin committed
269
270


271
# Modeling
thomwolf's avatar
thomwolf committed
272
if is_torch_available():
273
274
275
276
277
278
279
    # Benchmarks
    from .benchmark.benchmark import PyTorchBenchmark
    from .benchmark.benchmark_args import PyTorchBenchmarkArguments
    from .data.data_collator import (
        DataCollator,
        DataCollatorForLanguageModeling,
        DataCollatorForPermutationLanguageModeling,
280
        DataCollatorForSOP,
281
        DataCollatorForTokenClassification,
282
        DataCollatorForWholeWordMask,
283
284
285
286
287
288
289
        DataCollatorWithPadding,
        default_data_collator,
    )
    from .data.datasets import (
        GlueDataset,
        GlueDataTrainingArguments,
        LineByLineTextDataset,
290
        LineByLineWithRefDataset,
291
        LineByLineWithSOPTextDataset,
292
293
294
        SquadDataset,
        SquadDataTrainingArguments,
        TextDataset,
295
        TextDatasetForNextSentencePrediction,
296
    )
297
298
299
300
301
302
303
304
305
306
307
308
309
    from .generation_beam_search import BeamScorer, BeamSearchScorer
    from .generation_logits_process import (
        LogitsProcessor,
        LogitsProcessorList,
        LogitsWarper,
        MinLengthLogitsProcessor,
        NoBadWordsLogitsProcessor,
        NoRepeatNGramLogitsProcessor,
        RepetitionPenaltyLogitsProcessor,
        TemperatureLogitsWarper,
        TopKLogitsWarper,
        TopPLogitsWarper,
    )
310
    from .generation_utils import top_k_top_p_filtering
Sylvain Gugger's avatar
Sylvain Gugger committed
311
312
    from .modeling_utils import Conv1D, PreTrainedModel, apply_chunking_to_forward, prune_layer
    from .models.albert import (
313
314
315
316
317
318
319
320
321
322
323
        ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        AlbertForMaskedLM,
        AlbertForMultipleChoice,
        AlbertForPreTraining,
        AlbertForQuestionAnswering,
        AlbertForSequenceClassification,
        AlbertForTokenClassification,
        AlbertModel,
        AlbertPreTrainedModel,
        load_tf_weights_in_albert,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
324
    from .models.auto import (
325
326
        MODEL_FOR_CAUSAL_LM_MAPPING,
        MODEL_FOR_MASKED_LM_MAPPING,
327
        MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
328
        MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING,
329
330
        MODEL_FOR_PRETRAINING_MAPPING,
        MODEL_FOR_QUESTION_ANSWERING_MAPPING,
331
        MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
332
333
        MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
334
335
336
337
338
339
        MODEL_MAPPING,
        MODEL_WITH_LM_HEAD_MAPPING,
        AutoModel,
        AutoModelForCausalLM,
        AutoModelForMaskedLM,
        AutoModelForMultipleChoice,
340
        AutoModelForNextSentencePrediction,
341
342
343
344
345
346
        AutoModelForPreTraining,
        AutoModelForQuestionAnswering,
        AutoModelForSeq2SeqLM,
        AutoModelForSequenceClassification,
        AutoModelForTokenClassification,
        AutoModelWithLMHead,
347
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
348
    from .models.bart import (
349
350
351
352
353
354
        BART_PRETRAINED_MODEL_ARCHIVE_LIST,
        BartForConditionalGeneration,
        BartForQuestionAnswering,
        BartForSequenceClassification,
        BartModel,
        PretrainedBartModel,
Vasily Shamporov's avatar
Vasily Shamporov committed
355
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
356
    from .models.bert import (
357
        BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
358
        BertForMaskedLM,
359
        BertForMultipleChoice,
360
        BertForNextSentencePrediction,
361
362
        BertForPreTraining,
        BertForQuestionAnswering,
363
364
        BertForSequenceClassification,
        BertForTokenClassification,
Patrick von Platen's avatar
Patrick von Platen committed
365
        BertLayer,
366
367
368
369
        BertLMHeadModel,
        BertModel,
        BertPreTrainedModel,
        load_tf_weights_in_bert,
370
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
371
    from .models.bert_generation import (
372
373
374
375
        BertGenerationDecoder,
        BertGenerationEncoder,
        load_tf_weights_in_bert_generation,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
376
377
    from .models.blenderbot import BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST, BlenderbotForConditionalGeneration
    from .models.camembert import (
378
379
380
381
382
383
384
385
        CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        CamembertForCausalLM,
        CamembertForMaskedLM,
        CamembertForMultipleChoice,
        CamembertForQuestionAnswering,
        CamembertForSequenceClassification,
        CamembertForTokenClassification,
        CamembertModel,
386
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
387
388
    from .models.ctrl import CTRL_PRETRAINED_MODEL_ARCHIVE_LIST, CTRLLMHeadModel, CTRLModel, CTRLPreTrainedModel
    from .models.deberta import (
Pengcheng He's avatar
Pengcheng He committed
389
390
391
392
393
        DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
        DebertaForSequenceClassification,
        DebertaModel,
        DebertaPreTrainedModel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
394
    from .models.distilbert import (
395
396
397
398
399
400
401
402
403
        DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        DistilBertForMaskedLM,
        DistilBertForMultipleChoice,
        DistilBertForQuestionAnswering,
        DistilBertForSequenceClassification,
        DistilBertForTokenClassification,
        DistilBertModel,
        DistilBertPreTrainedModel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
404
    from .models.dpr import (
Ratthachat (Jung)'s avatar
Ratthachat (Jung) committed
405
406
407
        DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
408
409
410
411
412
413
414
        DPRContextEncoder,
        DPRPretrainedContextEncoder,
        DPRPretrainedQuestionEncoder,
        DPRPretrainedReader,
        DPRQuestionEncoder,
        DPRReader,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
415
    from .models.electra import (
416
417
418
419
420
421
422
423
424
425
426
        ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
        ElectraForMaskedLM,
        ElectraForMultipleChoice,
        ElectraForPreTraining,
        ElectraForQuestionAnswering,
        ElectraForSequenceClassification,
        ElectraForTokenClassification,
        ElectraModel,
        ElectraPreTrainedModel,
        load_tf_weights_in_electra,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
427
428
    from .models.encoder_decoder import EncoderDecoderModel
    from .models.flaubert import (
429
430
431
432
433
434
435
436
        FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        FlaubertForMultipleChoice,
        FlaubertForQuestionAnswering,
        FlaubertForQuestionAnsweringSimple,
        FlaubertForSequenceClassification,
        FlaubertForTokenClassification,
        FlaubertModel,
        FlaubertWithLMHeadModel,
437
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
438
439
    from .models.fsmt import FSMTForConditionalGeneration, FSMTModel, PretrainedFSMTModel
    from .models.funnel import (
Sylvain Gugger's avatar
Sylvain Gugger committed
440
441
442
443
444
445
446
447
448
449
450
        FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
        FunnelBaseModel,
        FunnelForMaskedLM,
        FunnelForMultipleChoice,
        FunnelForPreTraining,
        FunnelForQuestionAnswering,
        FunnelForSequenceClassification,
        FunnelForTokenClassification,
        FunnelModel,
        load_tf_weights_in_funnel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
451
    from .models.gpt2 import (
452
        GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
453
        GPT2DoubleHeadsModel,
454
        GPT2ForSequenceClassification,
455
456
457
        GPT2LMHeadModel,
        GPT2Model,
        GPT2PreTrainedModel,
458
459
        load_tf_weights_in_gpt2,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
460
    from .models.layoutlm import (
Minghao Li's avatar
Minghao Li committed
461
462
463
464
465
        LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        LayoutLMForMaskedLM,
        LayoutLMForTokenClassification,
        LayoutLMModel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
466
    from .models.longformer import (
467
468
469
470
471
472
473
474
        LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        LongformerForMaskedLM,
        LongformerForMultipleChoice,
        LongformerForQuestionAnswering,
        LongformerForSequenceClassification,
        LongformerForTokenClassification,
        LongformerModel,
        LongformerSelfAttention,
475
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
476
    from .models.lxmert import (
477
478
479
480
481
482
483
484
        LxmertEncoder,
        LxmertForPreTraining,
        LxmertForQuestionAnswering,
        LxmertModel,
        LxmertPreTrainedModel,
        LxmertVisualFeatureEncoder,
        LxmertXLayer,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
485
486
487
488
    from .models.marian import MarianMTModel
    from .models.mbart import MBartForConditionalGeneration
    from .models.mmbt import MMBTForClassification, MMBTModel, ModalEmbeddings
    from .models.mobilebert import (
489
490
491
492
493
494
495
496
497
498
499
500
501
        MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        MobileBertForMaskedLM,
        MobileBertForMultipleChoice,
        MobileBertForNextSentencePrediction,
        MobileBertForPreTraining,
        MobileBertForQuestionAnswering,
        MobileBertForSequenceClassification,
        MobileBertForTokenClassification,
        MobileBertLayer,
        MobileBertModel,
        MobileBertPreTrainedModel,
        load_tf_weights_in_mobilebert,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
502
    from .models.mt5 import MT5ForConditionalGeneration, MT5Model
Sylvain Gugger's avatar
Sylvain Gugger committed
503
    from .models.openai import (
504
505
        OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
        OpenAIGPTDoubleHeadsModel,
506
        OpenAIGPTForSequenceClassification,
507
508
509
510
        OpenAIGPTLMHeadModel,
        OpenAIGPTModel,
        OpenAIGPTPreTrainedModel,
        load_tf_weights_in_openai_gpt,
511
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
512
513
    from .models.pegasus import PegasusForConditionalGeneration
    from .models.prophetnet import (
Weizhen's avatar
Weizhen committed
514
515
516
517
518
519
520
521
        PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        ProphetNetDecoder,
        ProphetNetEncoder,
        ProphetNetForCausalLM,
        ProphetNetForConditionalGeneration,
        ProphetNetModel,
        ProphetNetPreTrainedModel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
522
523
    from .models.rag import RagModel, RagSequenceForGeneration, RagTokenForGeneration
    from .models.reformer import (
524
525
526
527
528
529
530
531
        REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        ReformerAttention,
        ReformerForMaskedLM,
        ReformerForQuestionAnswering,
        ReformerForSequenceClassification,
        ReformerLayer,
        ReformerModel,
        ReformerModelWithLMHead,
532
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
533
534
    from .models.retribert import RETRIBERT_PRETRAINED_MODEL_ARCHIVE_LIST, RetriBertModel, RetriBertPreTrainedModel
    from .models.roberta import (
535
        ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
536
        RobertaForCausalLM,
537
        RobertaForMaskedLM,
538
539
        RobertaForMultipleChoice,
        RobertaForQuestionAnswering,
540
541
542
        RobertaForSequenceClassification,
        RobertaForTokenClassification,
        RobertaModel,
543
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
544
    from .models.squeezebert import (
545
546
547
548
549
550
551
552
553
554
        SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        SqueezeBertForMaskedLM,
        SqueezeBertForMultipleChoice,
        SqueezeBertForQuestionAnswering,
        SqueezeBertForSequenceClassification,
        SqueezeBertForTokenClassification,
        SqueezeBertModel,
        SqueezeBertModule,
        SqueezeBertPreTrainedModel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
555
    from .models.t5 import (
556
        T5_PRETRAINED_MODEL_ARCHIVE_LIST,
557
        T5ForConditionalGeneration,
558
559
        T5Model,
        T5PreTrainedModel,
560
561
        load_tf_weights_in_t5,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
562
    from .models.transfo_xl import (
563
564
565
566
567
568
569
        TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
        AdaptiveEmbedding,
        TransfoXLLMHeadModel,
        TransfoXLModel,
        TransfoXLPreTrainedModel,
        load_tf_weights_in_transfo_xl,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
570
    from .models.xlm import (
571
572
573
574
575
576
577
578
579
        XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        XLMForMultipleChoice,
        XLMForQuestionAnswering,
        XLMForQuestionAnsweringSimple,
        XLMForSequenceClassification,
        XLMForTokenClassification,
        XLMModel,
        XLMPreTrainedModel,
        XLMWithLMHeadModel,
580
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
581
    from .models.xlm_prophetnet import (
Weizhen's avatar
Weizhen committed
582
583
584
585
586
587
588
        XLM_PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        XLMProphetNetDecoder,
        XLMProphetNetEncoder,
        XLMProphetNetForCausalLM,
        XLMProphetNetForConditionalGeneration,
        XLMProphetNetModel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
589
    from .models.xlm_roberta import (
590
        XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
591
        XLMRobertaForCausalLM,
592
593
        XLMRobertaForMaskedLM,
        XLMRobertaForMultipleChoice,
594
        XLMRobertaForQuestionAnswering,
595
596
        XLMRobertaForSequenceClassification,
        XLMRobertaForTokenClassification,
597
        XLMRobertaModel,
Quentin Lhoest's avatar
Quentin Lhoest committed
598
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
599
    from .models.xlnet import (
600
601
602
603
604
605
606
607
608
609
        XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        XLNetForMultipleChoice,
        XLNetForQuestionAnswering,
        XLNetForQuestionAnsweringSimple,
        XLNetForSequenceClassification,
        XLNetForTokenClassification,
        XLNetLMHeadModel,
        XLNetModel,
        XLNetPreTrainedModel,
        load_tf_weights_in_xlnet,
Yacine Jernite's avatar
Yacine Jernite committed
610
611
    )

thomwolf's avatar
thomwolf committed
612
    # Optimization
613
    from .optimization import (
614
        Adafactor,
615
616
617
618
619
620
        AdamW,
        get_constant_schedule,
        get_constant_schedule_with_warmup,
        get_cosine_schedule_with_warmup,
        get_cosine_with_hard_restarts_schedule_with_warmup,
        get_linear_schedule_with_warmup,
621
        get_polynomial_decay_schedule_with_warmup,
622
    )
thomwolf's avatar
thomwolf committed
623

Julien Chaumond's avatar
Julien Chaumond committed
624
    # Trainer
Sylvain Gugger's avatar
Sylvain Gugger committed
625
626
    from .trainer import Trainer
    from .trainer_pt_utils import torch_distributed_zero_first
627
628
else:
    from .utils.dummy_pt_objects import *
629

thomwolf's avatar
thomwolf committed
630
# TensorFlow
thomwolf's avatar
thomwolf committed
631
if is_tf_available():
632
633
634
635
    from .benchmark.benchmark_args_tf import TensorFlowBenchmarkArguments

    # Benchmarks
    from .benchmark.benchmark_tf import TensorFlowBenchmark
636
    from .generation_tf_utils import tf_top_k_top_p_filtering
Sylvain Gugger's avatar
Sylvain Gugger committed
637
638
    from .modeling_tf_utils import TFPreTrainedModel, TFSequenceSummary, TFSharedEmbeddings, shape_list
    from .models.albert import (
639
640
641
642
643
644
645
646
647
648
        TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFAlbertForMaskedLM,
        TFAlbertForMultipleChoice,
        TFAlbertForPreTraining,
        TFAlbertForQuestionAnswering,
        TFAlbertForSequenceClassification,
        TFAlbertForTokenClassification,
        TFAlbertMainLayer,
        TFAlbertModel,
        TFAlbertPreTrainedModel,
649
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
650
    from .models.auto import (
651
652
        TF_MODEL_FOR_CAUSAL_LM_MAPPING,
        TF_MODEL_FOR_MASKED_LM_MAPPING,
Julien Plu's avatar
Julien Plu committed
653
        TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
Julien Plu's avatar
Julien Plu committed
654
        TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING,
655
656
        TF_MODEL_FOR_PRETRAINING_MAPPING,
        TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
657
        TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
Julien Plu's avatar
Julien Plu committed
658
        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
659
        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
660
        TF_MODEL_MAPPING,
Julien Plu's avatar
Julien Plu committed
661
662
        TF_MODEL_WITH_LM_HEAD_MAPPING,
        TFAutoModel,
663
664
        TFAutoModelForCausalLM,
        TFAutoModelForMaskedLM,
Julien Plu's avatar
Julien Plu committed
665
666
667
        TFAutoModelForMultipleChoice,
        TFAutoModelForPreTraining,
        TFAutoModelForQuestionAnswering,
668
        TFAutoModelForSeq2SeqLM,
Julien Plu's avatar
Julien Plu committed
669
670
671
        TFAutoModelForSequenceClassification,
        TFAutoModelForTokenClassification,
        TFAutoModelWithLMHead,
672
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
673
674
    from .models.bart import TFBartForConditionalGeneration, TFBartModel
    from .models.bert import (
Julien Plu's avatar
Julien Plu committed
675
        TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
676
677
        TFBertEmbeddings,
        TFBertForMaskedLM,
Julien Plu's avatar
Julien Plu committed
678
        TFBertForMultipleChoice,
679
        TFBertForNextSentencePrediction,
Julien Plu's avatar
Julien Plu committed
680
681
        TFBertForPreTraining,
        TFBertForQuestionAnswering,
682
683
        TFBertForSequenceClassification,
        TFBertForTokenClassification,
684
        TFBertLMHeadModel,
Julien Plu's avatar
Julien Plu committed
685
686
687
        TFBertMainLayer,
        TFBertModel,
        TFBertPreTrainedModel,
688
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
689
690
    from .models.blenderbot import TFBlenderbotForConditionalGeneration
    from .models.camembert import (
Julien Plu's avatar
Julien Plu committed
691
692
693
694
695
696
        TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFCamembertForMaskedLM,
        TFCamembertForMultipleChoice,
        TFCamembertForQuestionAnswering,
        TFCamembertForSequenceClassification,
        TFCamembertForTokenClassification,
697
        TFCamembertModel,
698
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
699
    from .models.ctrl import (
Julien Plu's avatar
Julien Plu committed
700
701
702
703
        TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFCTRLLMHeadModel,
        TFCTRLModel,
        TFCTRLPreTrainedModel,
704
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
705
    from .models.distilbert import (
Julien Plu's avatar
Julien Plu committed
706
707
708
709
710
711
712
713
714
        TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDistilBertForMaskedLM,
        TFDistilBertForMultipleChoice,
        TFDistilBertForQuestionAnswering,
        TFDistilBertForSequenceClassification,
        TFDistilBertForTokenClassification,
        TFDistilBertMainLayer,
        TFDistilBertModel,
        TFDistilBertPreTrainedModel,
715
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
716
    from .models.dpr import (
Ratthachat (Jung)'s avatar
Ratthachat (Jung) committed
717
718
719
720
721
722
723
724
725
726
        TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDPRContextEncoder,
        TFDPRPretrainedContextEncoder,
        TFDPRPretrainedQuestionEncoder,
        TFDPRPretrainedReader,
        TFDPRQuestionEncoder,
        TFDPRReader,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
727
    from .models.electra import (
Julien Plu's avatar
Julien Plu committed
728
729
        TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFElectraForMaskedLM,
730
        TFElectraForMultipleChoice,
Julien Plu's avatar
Julien Plu committed
731
732
        TFElectraForPreTraining,
        TFElectraForQuestionAnswering,
733
        TFElectraForSequenceClassification,
Julien Plu's avatar
Julien Plu committed
734
735
736
        TFElectraForTokenClassification,
        TFElectraModel,
        TFElectraPreTrainedModel,
737
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
738
    from .models.flaubert import (
Julien Plu's avatar
Julien Plu committed
739
740
741
742
743
744
        TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFFlaubertForMultipleChoice,
        TFFlaubertForQuestionAnsweringSimple,
        TFFlaubertForSequenceClassification,
        TFFlaubertForTokenClassification,
        TFFlaubertModel,
745
        TFFlaubertWithLMHeadModel,
746
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
747
    from .models.funnel import (
748
749
750
751
752
753
754
755
756
757
        TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFFunnelBaseModel,
        TFFunnelForMaskedLM,
        TFFunnelForMultipleChoice,
        TFFunnelForPreTraining,
        TFFunnelForQuestionAnswering,
        TFFunnelForSequenceClassification,
        TFFunnelForTokenClassification,
        TFFunnelModel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
758
    from .models.gpt2 import (
Julien Plu's avatar
Julien Plu committed
759
760
761
762
763
764
765
        TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFGPT2DoubleHeadsModel,
        TFGPT2LMHeadModel,
        TFGPT2MainLayer,
        TFGPT2Model,
        TFGPT2PreTrainedModel,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
766
    from .models.longformer import (
Patrick von Platen's avatar
Patrick von Platen committed
767
768
769
        TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFLongformerForMaskedLM,
        TFLongformerForQuestionAnswering,
770
        TFLongformerModel,
Patrick von Platen's avatar
Patrick von Platen committed
771
772
        TFLongformerSelfAttention,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
773
    from .models.lxmert import (
774
775
776
777
778
779
780
        TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFLxmertForPreTraining,
        TFLxmertMainLayer,
        TFLxmertModel,
        TFLxmertPreTrainedModel,
        TFLxmertVisualFeatureEncoder,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
781
782
783
    from .models.marian import TFMarianMTModel
    from .models.mbart import TFMBartForConditionalGeneration
    from .models.mobilebert import (
Vasily Shamporov's avatar
Vasily Shamporov committed
784
785
786
        TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFMobileBertForMaskedLM,
        TFMobileBertForMultipleChoice,
787
788
789
790
        TFMobileBertForNextSentencePrediction,
        TFMobileBertForPreTraining,
        TFMobileBertForQuestionAnswering,
        TFMobileBertForSequenceClassification,
Vasily Shamporov's avatar
Vasily Shamporov committed
791
792
        TFMobileBertForTokenClassification,
        TFMobileBertMainLayer,
793
794
        TFMobileBertModel,
        TFMobileBertPreTrainedModel,
Vasily Shamporov's avatar
Vasily Shamporov committed
795
    )
Patrick von Platen's avatar
Patrick von Platen committed
796
    from .models.mt5 import TFMT5ForConditionalGeneration, TFMT5Model
Sylvain Gugger's avatar
Sylvain Gugger committed
797
    from .models.openai import (
Julien Plu's avatar
Julien Plu committed
798
799
800
801
802
803
        TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFOpenAIGPTDoubleHeadsModel,
        TFOpenAIGPTLMHeadModel,
        TFOpenAIGPTMainLayer,
        TFOpenAIGPTModel,
        TFOpenAIGPTPreTrainedModel,
Julien Plu's avatar
Julien Plu committed
804
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
805
806
    from .models.pegasus import TFPegasusForConditionalGeneration
    from .models.roberta import (
Julien Plu's avatar
Julien Plu committed
807
        TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
808
        TFRobertaForMaskedLM,
Julien Plu's avatar
Julien Plu committed
809
810
        TFRobertaForMultipleChoice,
        TFRobertaForQuestionAnswering,
811
812
        TFRobertaForSequenceClassification,
        TFRobertaForTokenClassification,
Julien Plu's avatar
Julien Plu committed
813
814
815
        TFRobertaMainLayer,
        TFRobertaModel,
        TFRobertaPreTrainedModel,
816
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
817
    from .models.t5 import (
Julien Plu's avatar
Julien Plu committed
818
819
820
821
        TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFT5ForConditionalGeneration,
        TFT5Model,
        TFT5PreTrainedModel,
822
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
823
    from .models.transfo_xl import (
Julien Plu's avatar
Julien Plu committed
824
825
826
827
828
829
        TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFAdaptiveEmbedding,
        TFTransfoXLLMHeadModel,
        TFTransfoXLMainLayer,
        TFTransfoXLModel,
        TFTransfoXLPreTrainedModel,
830
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
831
    from .models.xlm import (
Julien Plu's avatar
Julien Plu committed
832
833
834
835
836
837
838
839
        TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLMForMultipleChoice,
        TFXLMForQuestionAnsweringSimple,
        TFXLMForSequenceClassification,
        TFXLMForTokenClassification,
        TFXLMMainLayer,
        TFXLMModel,
        TFXLMPreTrainedModel,
840
        TFXLMWithLMHeadModel,
841
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
842
    from .models.xlm_roberta import (
Julien Plu's avatar
Julien Plu committed
843
844
845
846
847
848
849
        TF_XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLMRobertaForMaskedLM,
        TFXLMRobertaForMultipleChoice,
        TFXLMRobertaForQuestionAnswering,
        TFXLMRobertaForSequenceClassification,
        TFXLMRobertaForTokenClassification,
        TFXLMRobertaModel,
Julien Plu's avatar
Julien Plu committed
850
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
851
    from .models.xlnet import (
Julien Plu's avatar
Julien Plu committed
852
853
854
855
856
857
858
859
860
        TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLNetForMultipleChoice,
        TFXLNetForQuestionAnsweringSimple,
        TFXLNetForSequenceClassification,
        TFXLNetForTokenClassification,
        TFXLNetLMHeadModel,
        TFXLNetMainLayer,
        TFXLNetModel,
        TFXLNetPreTrainedModel,
Lysandre Debut's avatar
Lysandre Debut committed
861
862
    )

863
    # Optimization
864
    from .optimization_tf import AdamWeightDecay, GradientAccumulator, WarmUp, create_optimizer
Lysandre's avatar
Lysandre committed
865

Julien Plu's avatar
Julien Plu committed
866
867
868
    # Trainer
    from .trainer_tf import TFTrainer

869
870
871
872
873
else:
    # Import the same objects as dummies to get them in the namespace.
    # They will raise an import error if the user tries to instantiate / use them.
    from .utils.dummy_tf_objects import *

874

875
if is_flax_available():
Sylvain Gugger's avatar
Sylvain Gugger committed
876
877
    from .models.bert import FlaxBertModel
    from .models.roberta import FlaxRobertaModel
878
879
880
881
882
else:
    # Import the same objects as dummies to get them in the namespace.
    # They will raise an import error if the user tries to instantiate / use them.
    from .utils.dummy_flax_objects import *

883

884
if not is_tf_available() and not is_torch_available():
885
886
887
888
889
    logger.warning(
        "Neither PyTorch nor TensorFlow >= 2.0 have been found."
        "Models won't be available and only tokenizers, configuration"
        "and file/data utilities can be used."
    )