Views
No views yet
1base_model = 'klue/roberta-base'
2base_tokenizer = 'klue/roberta-base'
3
4from transformers import RobertaTokenizer, RobertaForMaskedLM
5from transformers import AutoModel, AutoTokenizer
6model = RobertaForMaskedLM.from_pretrained(base_model)
7tokenizer = AutoTokenizer.from_pretrained(base_tokenizer)
8
9from transformers import LineByLineTextDataset
10dataset = LineByLineTextDataset(
11 tokenizer=tokenizer,
12 file_path=fpath_dataset,
13 block_size=512,
14)
15
16from transformers import DataCollatorForLanguageModeling
17data_collator = DataCollatorForLanguageModeling(
18 tokenizer=tokenizer, mlm=True, mlm_probability=0.15
19)
20
21from transformers import Trainer, TrainingArguments
22training_args = TrainingArguments(
23 output_dir=output_dir,
24 overwrite_output_dir=True,
25 num_train_epochs=5,
26 per_device_train_batch_size=18,
27 save_steps=100,
28 save_total_limit=2,
29 seed=1
30)
31
32trainer = Trainer(
33 model=model,
34 args=training_args,
35 data_collator=data_collator,
36 train_dataset=dataset
37)
38
39train_metrics = trainer.train()
40trainer.save_model(output_dir)
41trainer.push_to_hub()