File size: 1,461 Bytes
78ab398 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import json
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("ghosh-r/bangla-gpt2")
train_path = 'train.txt'
test_path = 'test.txt'
from transformers import TextDataset,DataCollatorForLanguageModeling
def load_dataset(train_path,test_path,tokenizer):
train_dataset = TextDataset(
tokenizer=tokenizer,
file_path=train_path,
block_size=128)
test_dataset = TextDataset(
tokenizer=tokenizer,
file_path=test_path,
block_size=128)
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer, mlm=False,
)
return train_dataset, test_dataset, data_collator
train_dataset,test_dataset,data_collator = load_dataset(train_path,test_path,tokenizer)
from transformers import Trainer, TrainingArguments, AutoModelWithLMHead
model = AutoModelWithLMHead.from_pretrained("ghosh-r/bangla-gpt2")
training_args = TrainingArguments(
output_dir="./bn-poets",
overwrite_output_dir=True,
num_train_epochs=3,
per_device_train_batch_size=32,
per_device_eval_batch_size=64,
eval_steps = 400,
save_steps=800,
warmup_steps=500,
prediction_loss_only=True,
)
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=train_dataset,
eval_dataset=test_dataset,
)
trainer.train()
trainer.save_model()
tokenizer.save_pretrained('./bn-poets') |