__init__.py 27.5 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.

Sylvain Gugger's avatar
Sylvain Gugger committed
5
__version__ = "3.3.1"
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,
)

Julien Chaumond's avatar
Julien Chaumond committed
30
# Configurations
Aymeric Augustin's avatar
Aymeric Augustin committed
31
from .configuration_albert import ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, AlbertConfig
32
from .configuration_auto import ALL_PRETRAINED_CONFIG_ARCHIVE_MAP, CONFIG_MAPPING, AutoConfig
33
from .configuration_bart import BartConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
34
from .configuration_bert import BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, BertConfig
35
from .configuration_bert_generation import BertGenerationConfig
Sam Shleifer's avatar
Sam Shleifer committed
36
from .configuration_blenderbot import BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP, BlenderbotConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
37
38
from .configuration_camembert import CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, CamembertConfig
from .configuration_ctrl import CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP, CTRLConfig
Pengcheng He's avatar
Pengcheng He committed
39
from .configuration_deberta import DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, DebertaConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
40
from .configuration_distilbert import DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, DistilBertConfig
Quentin Lhoest's avatar
Quentin Lhoest committed
41
from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
Lysandre Debut's avatar
Lysandre Debut committed
42
from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig
43
from .configuration_encoder_decoder import EncoderDecoderConfig
Hang Le's avatar
Hang Le committed
44
from .configuration_flaubert import FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, FlaubertConfig
45
from .configuration_fsmt import FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP, FSMTConfig
Sylvain Gugger's avatar
Sylvain Gugger committed
46
from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
47
from .configuration_gpt2 import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPT2Config
Minghao Li's avatar
Minghao Li committed
48
from .configuration_layoutlm import LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMConfig
Iz Beltagy's avatar
Iz Beltagy committed
49
from .configuration_longformer import LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, LongformerConfig
50
from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
51
from .configuration_marian import MarianConfig
52
from .configuration_mbart import MBartConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
53
from .configuration_mmbt import MMBTConfig
Vasily Shamporov's avatar
Vasily Shamporov committed
54
from .configuration_mobilebert import MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MobileBertConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
55
from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
56
from .configuration_pegasus import PegasusConfig
Ola Piktus's avatar
Ola Piktus committed
57
from .configuration_rag import RagConfig
Patrick von Platen's avatar
Patrick von Platen committed
58
from .configuration_reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig
Yacine Jernite's avatar
Yacine Jernite committed
59
from .configuration_retribert import RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RetriBertConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
60
from .configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig
61
from .configuration_squeezebert import SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, SqueezeBertConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
62
63
64
65
66
67
from .configuration_t5 import T5_PRETRAINED_CONFIG_ARCHIVE_MAP, T5Config
from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
from .configuration_utils import PretrainedConfig
from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig
from .configuration_xlm_roberta import XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMRobertaConfig
from .configuration_xlnet import XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLNetConfig
68
from .data import (
Aymeric Augustin's avatar
Aymeric Augustin committed
69
    DataProcessor,
70
71
72
    InputExample,
    InputFeatures,
    SingleSentenceClassificationProcessor,
Aymeric Augustin's avatar
Aymeric Augustin committed
73
74
75
76
    SquadExample,
    SquadFeatures,
    SquadV1Processor,
    SquadV2Processor,
77
    glue_compute_metrics,
78
    glue_convert_examples_to_features,
Aymeric Augustin's avatar
Aymeric Augustin committed
79
    glue_output_modes,
80
81
    glue_processors,
    glue_tasks_num_labels,
Aymeric Augustin's avatar
Aymeric Augustin committed
82
    squad_convert_examples_to_features,
83
    xnli_compute_metrics,
84
85
86
87
    xnli_output_modes,
    xnli_processors,
    xnli_tasks_num_labels,
)
88

Aymeric Augustin's avatar
Aymeric Augustin committed
89
90
91
92
93
94
95
96
97
98
99
100
101
# Files and general utilities
from .file_utils import (
    CONFIG_NAME,
    MODEL_CARD_NAME,
    PYTORCH_PRETRAINED_BERT_CACHE,
    PYTORCH_TRANSFORMERS_CACHE,
    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
102
    is_apex_available,
103
    is_datasets_available,
Ola Piktus's avatar
Ola Piktus committed
104
    is_faiss_available,
Patrick von Platen's avatar
Patrick von Platen committed
105
106
    is_psutil_available,
    is_py3nvml_available,
107
    is_sklearn_available,
Aymeric Augustin's avatar
Aymeric Augustin committed
108
109
    is_tf_available,
    is_torch_available,
110
    is_torch_tpu_available,
Aymeric Augustin's avatar
Aymeric Augustin committed
111
)
112
from .hf_argparser import HfArgumentParser
113

thomwolf's avatar
thomwolf committed
114
# Model Cards
115
from .modelcard import ModelCard
116

Aymeric Augustin's avatar
Aymeric Augustin committed
117
118
119
120
121
122
123
124
125
126
# 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,
)
127

Aymeric Augustin's avatar
Aymeric Augustin committed
128
129
# Pipelines
from .pipelines import (
130
131
    Conversation,
    ConversationalPipeline,
Aymeric Augustin's avatar
Aymeric Augustin committed
132
133
    CsvPipelineDataFormat,
    FeatureExtractionPipeline,
Julien Chaumond's avatar
Julien Chaumond committed
134
    FillMaskPipeline,
Aymeric Augustin's avatar
Aymeric Augustin committed
135
136
137
138
139
140
    JsonPipelineDataFormat,
    NerPipeline,
    PipedPipelineDataFormat,
    Pipeline,
    PipelineDataFormat,
    QuestionAnsweringPipeline,
141
    SummarizationPipeline,
142
    Text2TextGenerationPipeline,
Aymeric Augustin's avatar
Aymeric Augustin committed
143
    TextClassificationPipeline,
144
    TextGenerationPipeline,
145
    TokenClassificationPipeline,
146
    TranslationPipeline,
147
    ZeroShotClassificationPipeline,
Aymeric Augustin's avatar
Aymeric Augustin committed
148
149
    pipeline,
)
Julien Chaumond's avatar
Julien Chaumond committed
150

Ola Piktus's avatar
Ola Piktus committed
151
152
153
# Retriever
from .retrieval_rag import RagRetriever

Julien Chaumond's avatar
Julien Chaumond committed
154
# Tokenizers
155
from .tokenization_albert import AlbertTokenizer, AlbertTokenizerFast
156
from .tokenization_auto import TOKENIZER_MAPPING, AutoTokenizer
157
from .tokenization_bart import BartTokenizer, BartTokenizerFast
Anthony MOI's avatar
Anthony MOI committed
158
from .tokenization_bert import BasicTokenizer, BertTokenizer, BertTokenizerFast, WordpieceTokenizer
159
from .tokenization_bert_generation import BertGenerationTokenizer
Aymeric Augustin's avatar
Aymeric Augustin committed
160
from .tokenization_bert_japanese import BertJapaneseTokenizer, CharacterTokenizer, MecabTokenizer
161
from .tokenization_bertweet import BertweetTokenizer
Sam Shleifer's avatar
Sam Shleifer committed
162
from .tokenization_blenderbot import BlenderbotSmallTokenizer, BlenderbotTokenizer
163
from .tokenization_camembert import CamembertTokenizer, CamembertTokenizerFast
keskarnitish's avatar
keskarnitish committed
164
from .tokenization_ctrl import CTRLTokenizer
Pengcheng He's avatar
Pengcheng He committed
165
from .tokenization_deberta import DebertaTokenizer
166
from .tokenization_distilbert import DistilBertTokenizer, DistilBertTokenizerFast
Quentin Lhoest's avatar
Quentin Lhoest committed
167
168
169
170
171
172
173
174
from .tokenization_dpr import (
    DPRContextEncoderTokenizer,
    DPRContextEncoderTokenizerFast,
    DPRQuestionEncoderTokenizer,
    DPRQuestionEncoderTokenizerFast,
    DPRReaderTokenizer,
    DPRReaderTokenizerFast,
)
Lysandre Debut's avatar
Lysandre Debut committed
175
from .tokenization_electra import ElectraTokenizer, ElectraTokenizerFast
Hang Le's avatar
Hang Le committed
176
from .tokenization_flaubert import FlaubertTokenizer
177
from .tokenization_fsmt import FSMTTokenizer
Sylvain Gugger's avatar
Sylvain Gugger committed
178
from .tokenization_funnel import FunnelTokenizer, FunnelTokenizerFast
Anthony MOI's avatar
Anthony MOI committed
179
from .tokenization_gpt2 import GPT2Tokenizer, GPT2TokenizerFast
rmroczkowski's avatar
rmroczkowski committed
180
from .tokenization_herbert import HerbertTokenizer, HerbertTokenizerFast
Minghao Li's avatar
Minghao Li committed
181
from .tokenization_layoutlm import LayoutLMTokenizer, LayoutLMTokenizerFast
Suraj Patil's avatar
Suraj Patil committed
182
from .tokenization_longformer import LongformerTokenizer, LongformerTokenizerFast
183
from .tokenization_lxmert import LxmertTokenizer, LxmertTokenizerFast
184
from .tokenization_mbart import MBartTokenizer, MBartTokenizerFast
Vasily Shamporov's avatar
Vasily Shamporov committed
185
from .tokenization_mobilebert import MobileBertTokenizer, MobileBertTokenizerFast
186
from .tokenization_openai import OpenAIGPTTokenizer, OpenAIGPTTokenizerFast
187
from .tokenization_pegasus import PegasusTokenizer, PegasusTokenizerFast
188
from .tokenization_phobert import PhobertTokenizer
Ola Piktus's avatar
Ola Piktus committed
189
from .tokenization_rag import RagTokenizer
190
from .tokenization_reformer import ReformerTokenizer, ReformerTokenizerFast
Yacine Jernite's avatar
Yacine Jernite committed
191
from .tokenization_retribert import RetriBertTokenizer, RetriBertTokenizerFast
192
from .tokenization_roberta import RobertaTokenizer, RobertaTokenizerFast
193
from .tokenization_squeezebert import SqueezeBertTokenizer, SqueezeBertTokenizerFast
194
195
from .tokenization_t5 import T5Tokenizer, T5TokenizerFast
from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer
196
197
from .tokenization_utils import PreTrainedTokenizer
from .tokenization_utils_base import (
198
    BatchEncoding,
199
200
    CharSpan,
    PreTrainedTokenizerBase,
201
202
    SpecialTokensMixin,
    TensorType,
203
    TokenSpan,
204
)
205
from .tokenization_utils_fast import PreTrainedTokenizerFast
Aymeric Augustin's avatar
Aymeric Augustin committed
206
from .tokenization_xlm import XLMTokenizer
207
208
from .tokenization_xlm_roberta import XLMRobertaTokenizer, XLMRobertaTokenizerFast
from .tokenization_xlnet import SPIECE_UNDERLINE, XLNetTokenizer, XLNetTokenizerFast
209
210

# Trainer
Sylvain Gugger's avatar
Sylvain Gugger committed
211
212
213
214
215
216
217
218
219
from .trainer_callback import (
    DefaultFlowCallback,
    PrinterCallback,
    ProgressCallback,
    TrainerCallback,
    TrainerControl,
    TrainerState,
)
from .trainer_utils import EvalPrediction, EvaluationStrategy, set_seed
220
from .training_args import TrainingArguments
Julien Plu's avatar
Julien Plu committed
221
from .training_args_tf import TFTrainingArguments
Lysandre Debut's avatar
Lysandre Debut committed
222
from .utils import logging
Aymeric Augustin's avatar
Aymeric Augustin committed
223
224


Lysandre Debut's avatar
Lysandre Debut committed
225
logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
Aymeric Augustin's avatar
Aymeric Augustin committed
226
227


228
# Modeling
thomwolf's avatar
thomwolf committed
229
if is_torch_available():
230
231
232
233
234
235
    # Benchmarks
    from .benchmark.benchmark import PyTorchBenchmark
    from .benchmark.benchmark_args import PyTorchBenchmarkArguments
    from .data.data_collator import (
        DataCollator,
        DataCollatorForLanguageModeling,
236
        DataCollatorForNextSentencePrediction,
237
        DataCollatorForPermutationLanguageModeling,
238
        DataCollatorForSOP,
239
240
241
242
243
244
245
        DataCollatorWithPadding,
        default_data_collator,
    )
    from .data.datasets import (
        GlueDataset,
        GlueDataTrainingArguments,
        LineByLineTextDataset,
246
        LineByLineWithSOPTextDataset,
247
248
249
        SquadDataset,
        SquadDataTrainingArguments,
        TextDataset,
250
        TextDatasetForNextSentencePrediction,
251
    )
252
    from .generation_utils import top_k_top_p_filtering
253
254
255
256
257
258
259
260
261
262
263
264
    from .modeling_albert import (
        ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        AlbertForMaskedLM,
        AlbertForMultipleChoice,
        AlbertForPreTraining,
        AlbertForQuestionAnswering,
        AlbertForSequenceClassification,
        AlbertForTokenClassification,
        AlbertModel,
        AlbertPreTrainedModel,
        load_tf_weights_in_albert,
    )
265
    from .modeling_auto import (
266
267
        MODEL_FOR_CAUSAL_LM_MAPPING,
        MODEL_FOR_MASKED_LM_MAPPING,
268
269
270
        MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
        MODEL_FOR_PRETRAINING_MAPPING,
        MODEL_FOR_QUESTION_ANSWERING_MAPPING,
271
        MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
272
273
        MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
274
275
276
277
278
279
280
281
282
283
284
285
        MODEL_MAPPING,
        MODEL_WITH_LM_HEAD_MAPPING,
        AutoModel,
        AutoModelForCausalLM,
        AutoModelForMaskedLM,
        AutoModelForMultipleChoice,
        AutoModelForPreTraining,
        AutoModelForQuestionAnswering,
        AutoModelForSeq2SeqLM,
        AutoModelForSequenceClassification,
        AutoModelForTokenClassification,
        AutoModelWithLMHead,
286
    )
287
288
289
290
291
292
293
    from .modeling_bart import (
        BART_PRETRAINED_MODEL_ARCHIVE_LIST,
        BartForConditionalGeneration,
        BartForQuestionAnswering,
        BartForSequenceClassification,
        BartModel,
        PretrainedBartModel,
Vasily Shamporov's avatar
Vasily Shamporov committed
294
    )
295
    from .modeling_bert import (
296
        BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
297
        BertForMaskedLM,
298
        BertForMultipleChoice,
299
        BertForNextSentencePrediction,
300
301
        BertForPreTraining,
        BertForQuestionAnswering,
302
303
        BertForSequenceClassification,
        BertForTokenClassification,
Patrick von Platen's avatar
Patrick von Platen committed
304
        BertLayer,
305
306
307
308
        BertLMHeadModel,
        BertModel,
        BertPreTrainedModel,
        load_tf_weights_in_bert,
309
    )
310
311
312
313
314
    from .modeling_bert_generation import (
        BertGenerationDecoder,
        BertGenerationEncoder,
        load_tf_weights_in_bert_generation,
    )
Sam Shleifer's avatar
Sam Shleifer committed
315
    from .modeling_blenderbot import BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST, BlenderbotForConditionalGeneration
316
317
318
319
320
321
322
323
324
    from .modeling_camembert import (
        CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        CamembertForCausalLM,
        CamembertForMaskedLM,
        CamembertForMultipleChoice,
        CamembertForQuestionAnswering,
        CamembertForSequenceClassification,
        CamembertForTokenClassification,
        CamembertModel,
325
    )
326
    from .modeling_ctrl import CTRL_PRETRAINED_MODEL_ARCHIVE_LIST, CTRLLMHeadModel, CTRLModel, CTRLPreTrainedModel
Pengcheng He's avatar
Pengcheng He committed
327
328
329
330
331
332
    from .modeling_deberta import (
        DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
        DebertaForSequenceClassification,
        DebertaModel,
        DebertaPreTrainedModel,
    )
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
    from .modeling_distilbert import (
        DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        DistilBertForMaskedLM,
        DistilBertForMultipleChoice,
        DistilBertForQuestionAnswering,
        DistilBertForSequenceClassification,
        DistilBertForTokenClassification,
        DistilBertModel,
        DistilBertPreTrainedModel,
    )
    from .modeling_dpr import (
        DPRContextEncoder,
        DPRPretrainedContextEncoder,
        DPRPretrainedQuestionEncoder,
        DPRPretrainedReader,
        DPRQuestionEncoder,
        DPRReader,
    )
    from .modeling_electra import (
        ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
        ElectraForMaskedLM,
        ElectraForMultipleChoice,
        ElectraForPreTraining,
        ElectraForQuestionAnswering,
        ElectraForSequenceClassification,
        ElectraForTokenClassification,
        ElectraModel,
        ElectraPreTrainedModel,
        load_tf_weights_in_electra,
    )
    from .modeling_encoder_decoder import EncoderDecoderModel
    from .modeling_flaubert import (
        FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        FlaubertForMultipleChoice,
        FlaubertForQuestionAnswering,
        FlaubertForQuestionAnsweringSimple,
        FlaubertForSequenceClassification,
        FlaubertForTokenClassification,
        FlaubertModel,
        FlaubertWithLMHeadModel,
373
    )
374
    from .modeling_fsmt import FSMTForConditionalGeneration, FSMTModel, PretrainedFSMTModel
Sylvain Gugger's avatar
Sylvain Gugger committed
375
376
377
378
379
380
381
382
383
384
385
386
    from .modeling_funnel import (
        FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
        FunnelBaseModel,
        FunnelForMaskedLM,
        FunnelForMultipleChoice,
        FunnelForPreTraining,
        FunnelForQuestionAnswering,
        FunnelForSequenceClassification,
        FunnelForTokenClassification,
        FunnelModel,
        load_tf_weights_in_funnel,
    )
387
    from .modeling_gpt2 import (
388
        GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
389
        GPT2DoubleHeadsModel,
390
        GPT2ForSequenceClassification,
391
392
393
        GPT2LMHeadModel,
        GPT2Model,
        GPT2PreTrainedModel,
394
395
        load_tf_weights_in_gpt2,
    )
Minghao Li's avatar
Minghao Li committed
396
397
398
399
400
401
    from .modeling_layoutlm import (
        LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        LayoutLMForMaskedLM,
        LayoutLMForTokenClassification,
        LayoutLMModel,
    )
402
403
404
405
406
407
408
409
410
    from .modeling_longformer import (
        LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        LongformerForMaskedLM,
        LongformerForMultipleChoice,
        LongformerForQuestionAnswering,
        LongformerForSequenceClassification,
        LongformerForTokenClassification,
        LongformerModel,
        LongformerSelfAttention,
411
    )
412
413
414
415
416
417
418
419
420
    from .modeling_lxmert import (
        LxmertEncoder,
        LxmertForPreTraining,
        LxmertForQuestionAnswering,
        LxmertModel,
        LxmertPreTrainedModel,
        LxmertVisualFeatureEncoder,
        LxmertXLayer,
    )
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
    from .modeling_marian import MarianMTModel
    from .modeling_mbart import MBartForConditionalGeneration
    from .modeling_mmbt import MMBTForClassification, MMBTModel, ModalEmbeddings
    from .modeling_mobilebert import (
        MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        MobileBertForMaskedLM,
        MobileBertForMultipleChoice,
        MobileBertForNextSentencePrediction,
        MobileBertForPreTraining,
        MobileBertForQuestionAnswering,
        MobileBertForSequenceClassification,
        MobileBertForTokenClassification,
        MobileBertLayer,
        MobileBertModel,
        MobileBertPreTrainedModel,
        load_tf_weights_in_mobilebert,
    )
    from .modeling_openai import (
        OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
        OpenAIGPTDoubleHeadsModel,
441
        OpenAIGPTForSequenceClassification,
442
443
444
445
        OpenAIGPTLMHeadModel,
        OpenAIGPTModel,
        OpenAIGPTPreTrainedModel,
        load_tf_weights_in_openai_gpt,
446
    )
447
    from .modeling_pegasus import PegasusForConditionalGeneration
Ola Piktus's avatar
Ola Piktus committed
448
    from .modeling_rag import RagModel, RagSequenceForGeneration, RagTokenForGeneration
449
450
451
452
453
454
455
456
457
    from .modeling_reformer import (
        REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        ReformerAttention,
        ReformerForMaskedLM,
        ReformerForQuestionAnswering,
        ReformerForSequenceClassification,
        ReformerLayer,
        ReformerModel,
        ReformerModelWithLMHead,
458
    )
459
    from .modeling_retribert import RETRIBERT_PRETRAINED_MODEL_ARCHIVE_LIST, RetriBertModel, RetriBertPreTrainedModel
460
    from .modeling_roberta import (
461
        ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
462
        RobertaForCausalLM,
463
        RobertaForMaskedLM,
464
465
        RobertaForMultipleChoice,
        RobertaForQuestionAnswering,
466
467
468
        RobertaForSequenceClassification,
        RobertaForTokenClassification,
        RobertaModel,
469
    )
470
471
472
473
474
475
476
477
478
479
480
    from .modeling_squeezebert import (
        SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        SqueezeBertForMaskedLM,
        SqueezeBertForMultipleChoice,
        SqueezeBertForQuestionAnswering,
        SqueezeBertForSequenceClassification,
        SqueezeBertForTokenClassification,
        SqueezeBertModel,
        SqueezeBertModule,
        SqueezeBertPreTrainedModel,
    )
481
    from .modeling_t5 import (
482
        T5_PRETRAINED_MODEL_ARCHIVE_LIST,
483
        T5ForConditionalGeneration,
484
485
        T5Model,
        T5PreTrainedModel,
486
487
        load_tf_weights_in_t5,
    )
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
    from .modeling_transfo_xl import (
        TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
        AdaptiveEmbedding,
        TransfoXLLMHeadModel,
        TransfoXLModel,
        TransfoXLPreTrainedModel,
        load_tf_weights_in_transfo_xl,
    )
    from .modeling_utils import Conv1D, PreTrainedModel, apply_chunking_to_forward, prune_layer
    from .modeling_xlm import (
        XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        XLMForMultipleChoice,
        XLMForQuestionAnswering,
        XLMForQuestionAnsweringSimple,
        XLMForSequenceClassification,
        XLMForTokenClassification,
        XLMModel,
        XLMPreTrainedModel,
        XLMWithLMHeadModel,
507
508
    )
    from .modeling_xlm_roberta import (
509
        XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
510
        XLMRobertaForCausalLM,
511
512
        XLMRobertaForMaskedLM,
        XLMRobertaForMultipleChoice,
513
        XLMRobertaForQuestionAnswering,
514
515
        XLMRobertaForSequenceClassification,
        XLMRobertaForTokenClassification,
516
        XLMRobertaModel,
Quentin Lhoest's avatar
Quentin Lhoest committed
517
    )
518
519
520
521
522
523
524
525
526
527
528
    from .modeling_xlnet import (
        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
529
530
    )

thomwolf's avatar
thomwolf committed
531
    # Optimization
532
    from .optimization import (
533
        Adafactor,
534
535
536
537
538
539
        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,
540
        get_polynomial_decay_schedule_with_warmup,
541
    )
542
    from .tokenization_marian import MarianTokenizer
thomwolf's avatar
thomwolf committed
543

Julien Chaumond's avatar
Julien Chaumond committed
544
    # Trainer
Sylvain Gugger's avatar
Sylvain Gugger committed
545
546
    from .trainer import Trainer
    from .trainer_pt_utils import torch_distributed_zero_first
547
548
else:
    from .utils.dummy_pt_objects import *
549

thomwolf's avatar
thomwolf committed
550
# TensorFlow
thomwolf's avatar
thomwolf committed
551
if is_tf_available():
552
553
554
555
    from .benchmark.benchmark_args_tf import TensorFlowBenchmarkArguments

    # Benchmarks
    from .benchmark.benchmark_tf import TensorFlowBenchmark
556
    from .generation_tf_utils import tf_top_k_top_p_filtering
557
558
559
560
561
562
563
564
565
566
567
    from .modeling_tf_albert import (
        TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFAlbertForMaskedLM,
        TFAlbertForMultipleChoice,
        TFAlbertForPreTraining,
        TFAlbertForQuestionAnswering,
        TFAlbertForSequenceClassification,
        TFAlbertForTokenClassification,
        TFAlbertMainLayer,
        TFAlbertModel,
        TFAlbertPreTrainedModel,
568
    )
569
    from .modeling_tf_auto import (
570
571
        TF_MODEL_FOR_CAUSAL_LM_MAPPING,
        TF_MODEL_FOR_MASKED_LM_MAPPING,
Julien Plu's avatar
Julien Plu committed
572
        TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
573
574
        TF_MODEL_FOR_PRETRAINING_MAPPING,
        TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
575
        TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
Julien Plu's avatar
Julien Plu committed
576
        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
577
        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
578
        TF_MODEL_MAPPING,
Julien Plu's avatar
Julien Plu committed
579
580
        TF_MODEL_WITH_LM_HEAD_MAPPING,
        TFAutoModel,
581
582
        TFAutoModelForCausalLM,
        TFAutoModelForMaskedLM,
Julien Plu's avatar
Julien Plu committed
583
584
585
        TFAutoModelForMultipleChoice,
        TFAutoModelForPreTraining,
        TFAutoModelForQuestionAnswering,
586
        TFAutoModelForSeq2SeqLM,
Julien Plu's avatar
Julien Plu committed
587
588
589
        TFAutoModelForSequenceClassification,
        TFAutoModelForTokenClassification,
        TFAutoModelWithLMHead,
590
591
    )
    from .modeling_tf_bert import (
Julien Plu's avatar
Julien Plu committed
592
        TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
593
594
        TFBertEmbeddings,
        TFBertForMaskedLM,
Julien Plu's avatar
Julien Plu committed
595
        TFBertForMultipleChoice,
596
        TFBertForNextSentencePrediction,
Julien Plu's avatar
Julien Plu committed
597
598
        TFBertForPreTraining,
        TFBertForQuestionAnswering,
599
600
        TFBertForSequenceClassification,
        TFBertForTokenClassification,
601
        TFBertLMHeadModel,
Julien Plu's avatar
Julien Plu committed
602
603
604
        TFBertMainLayer,
        TFBertModel,
        TFBertPreTrainedModel,
605
    )
Julien Plu's avatar
Julien Plu committed
606
607
608
609
610
611
612
    from .modeling_tf_camembert import (
        TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFCamembertForMaskedLM,
        TFCamembertForMultipleChoice,
        TFCamembertForQuestionAnswering,
        TFCamembertForSequenceClassification,
        TFCamembertForTokenClassification,
613
        TFCamembertModel,
614
    )
Julien Plu's avatar
Julien Plu committed
615
616
617
618
619
    from .modeling_tf_ctrl import (
        TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFCTRLLMHeadModel,
        TFCTRLModel,
        TFCTRLPreTrainedModel,
620
    )
Julien Plu's avatar
Julien Plu committed
621
622
623
624
625
626
627
628
629
630
    from .modeling_tf_distilbert import (
        TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDistilBertForMaskedLM,
        TFDistilBertForMultipleChoice,
        TFDistilBertForQuestionAnswering,
        TFDistilBertForSequenceClassification,
        TFDistilBertForTokenClassification,
        TFDistilBertMainLayer,
        TFDistilBertModel,
        TFDistilBertPreTrainedModel,
631
    )
Julien Plu's avatar
Julien Plu committed
632
633
634
    from .modeling_tf_electra import (
        TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFElectraForMaskedLM,
635
        TFElectraForMultipleChoice,
Julien Plu's avatar
Julien Plu committed
636
637
        TFElectraForPreTraining,
        TFElectraForQuestionAnswering,
638
        TFElectraForSequenceClassification,
Julien Plu's avatar
Julien Plu committed
639
640
641
        TFElectraForTokenClassification,
        TFElectraModel,
        TFElectraPreTrainedModel,
642
    )
Julien Plu's avatar
Julien Plu committed
643
644
645
646
647
648
649
    from .modeling_tf_flaubert import (
        TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFFlaubertForMultipleChoice,
        TFFlaubertForQuestionAnsweringSimple,
        TFFlaubertForSequenceClassification,
        TFFlaubertForTokenClassification,
        TFFlaubertModel,
650
        TFFlaubertWithLMHeadModel,
651
    )
652
653
654
655
656
657
658
659
660
661
662
    from .modeling_tf_funnel import (
        TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFFunnelBaseModel,
        TFFunnelForMaskedLM,
        TFFunnelForMultipleChoice,
        TFFunnelForPreTraining,
        TFFunnelForQuestionAnswering,
        TFFunnelForSequenceClassification,
        TFFunnelForTokenClassification,
        TFFunnelModel,
    )
Julien Plu's avatar
Julien Plu committed
663
664
665
666
667
668
669
670
    from .modeling_tf_gpt2 import (
        TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFGPT2DoubleHeadsModel,
        TFGPT2LMHeadModel,
        TFGPT2MainLayer,
        TFGPT2Model,
        TFGPT2PreTrainedModel,
    )
Patrick von Platen's avatar
Patrick von Platen committed
671
672
673
674
    from .modeling_tf_longformer import (
        TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFLongformerForMaskedLM,
        TFLongformerForQuestionAnswering,
675
        TFLongformerModel,
Patrick von Platen's avatar
Patrick von Platen committed
676
677
        TFLongformerSelfAttention,
    )
678
679
680
681
682
683
684
685
    from .modeling_tf_lxmert import (
        TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFLxmertForPreTraining,
        TFLxmertMainLayer,
        TFLxmertModel,
        TFLxmertPreTrainedModel,
        TFLxmertVisualFeatureEncoder,
    )
Vasily Shamporov's avatar
Vasily Shamporov committed
686
687
688
689
    from .modeling_tf_mobilebert import (
        TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFMobileBertForMaskedLM,
        TFMobileBertForMultipleChoice,
690
691
692
693
        TFMobileBertForNextSentencePrediction,
        TFMobileBertForPreTraining,
        TFMobileBertForQuestionAnswering,
        TFMobileBertForSequenceClassification,
Vasily Shamporov's avatar
Vasily Shamporov committed
694
695
        TFMobileBertForTokenClassification,
        TFMobileBertMainLayer,
696
697
        TFMobileBertModel,
        TFMobileBertPreTrainedModel,
Vasily Shamporov's avatar
Vasily Shamporov committed
698
    )
Julien Plu's avatar
Julien Plu committed
699
700
701
702
703
704
705
    from .modeling_tf_openai import (
        TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFOpenAIGPTDoubleHeadsModel,
        TFOpenAIGPTLMHeadModel,
        TFOpenAIGPTMainLayer,
        TFOpenAIGPTModel,
        TFOpenAIGPTPreTrainedModel,
Julien Plu's avatar
Julien Plu committed
706
    )
707
    from .modeling_tf_roberta import (
Julien Plu's avatar
Julien Plu committed
708
        TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
709
        TFRobertaForMaskedLM,
Julien Plu's avatar
Julien Plu committed
710
711
        TFRobertaForMultipleChoice,
        TFRobertaForQuestionAnswering,
712
713
        TFRobertaForSequenceClassification,
        TFRobertaForTokenClassification,
Julien Plu's avatar
Julien Plu committed
714
715
716
        TFRobertaMainLayer,
        TFRobertaModel,
        TFRobertaPreTrainedModel,
717
    )
Julien Plu's avatar
Julien Plu committed
718
719
720
721
722
    from .modeling_tf_t5 import (
        TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFT5ForConditionalGeneration,
        TFT5Model,
        TFT5PreTrainedModel,
723
    )
Julien Plu's avatar
Julien Plu committed
724
725
726
727
728
729
730
    from .modeling_tf_transfo_xl import (
        TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFAdaptiveEmbedding,
        TFTransfoXLLMHeadModel,
        TFTransfoXLMainLayer,
        TFTransfoXLModel,
        TFTransfoXLPreTrainedModel,
731
    )
732
    from .modeling_tf_utils import TFPreTrainedModel, TFSequenceSummary, TFSharedEmbeddings, shape_list
Julien Plu's avatar
Julien Plu committed
733
734
735
736
737
738
739
740
741
    from .modeling_tf_xlm import (
        TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLMForMultipleChoice,
        TFXLMForQuestionAnsweringSimple,
        TFXLMForSequenceClassification,
        TFXLMForTokenClassification,
        TFXLMMainLayer,
        TFXLMModel,
        TFXLMPreTrainedModel,
742
        TFXLMWithLMHeadModel,
743
    )
Julien Plu's avatar
Julien Plu committed
744
745
746
747
748
749
750
751
    from .modeling_tf_xlm_roberta import (
        TF_XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLMRobertaForMaskedLM,
        TFXLMRobertaForMultipleChoice,
        TFXLMRobertaForQuestionAnswering,
        TFXLMRobertaForSequenceClassification,
        TFXLMRobertaForTokenClassification,
        TFXLMRobertaModel,
Julien Plu's avatar
Julien Plu committed
752
    )
Julien Plu's avatar
Julien Plu committed
753
754
755
756
757
758
759
760
761
762
    from .modeling_tf_xlnet import (
        TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLNetForMultipleChoice,
        TFXLNetForQuestionAnsweringSimple,
        TFXLNetForSequenceClassification,
        TFXLNetForTokenClassification,
        TFXLNetLMHeadModel,
        TFXLNetMainLayer,
        TFXLNetModel,
        TFXLNetPreTrainedModel,
Lysandre Debut's avatar
Lysandre Debut committed
763
764
    )

765
    # Optimization
766
    from .optimization_tf import AdamWeightDecay, GradientAccumulator, WarmUp, create_optimizer
Lysandre's avatar
Lysandre committed
767

Julien Plu's avatar
Julien Plu committed
768
769
770
    # Trainer
    from .trainer_tf import TFTrainer

771
772
773
774
775
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 *

776

777
if not is_tf_available() and not is_torch_available():
778
779
780
781
782
    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."
    )