|
@@ -11,7 +11,7 @@ import transformers
|
|
|
from datasets import load_from_disk
|
|
|
from torch.utils.data import DataLoader
|
|
|
from torch_optimizer import Lamb
|
|
|
-from transformers import DataCollatorForLanguageModeling, HfArgumentParser, TrainingArguments, set_seed, AdamW
|
|
|
+from transformers import DataCollatorForLanguageModeling, HfArgumentParser, TrainingArguments, set_seed, Adafactor
|
|
|
from transformers.models.albert import AlbertConfig, AlbertForPreTraining, AlbertTokenizerFast
|
|
|
from transformers.optimization import get_linear_schedule_with_warmup
|
|
|
from transformers.trainer import Trainer
|