File size: 2,108 Bytes
0ed8f16 |
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 os
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from datasets import load_dataset
def load_model_and_tokenizer(model_name):
"""
Load the model and tokenizer.
"""
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
return model, tokenizer
def load_and_tokenize_dataset(dataset_name, tokenizer, max_length=512):
"""
Load and tokenize the dataset.
"""
dataset = load_dataset(dataset_name)
def tokenize_function(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=max_length)
tokenized_datasets = dataset.map(tokenize_function, batched=True)
return tokenized_datasets
def setup_training_args(output_dir="./results", per_device_train_batch_size=2, per_device_eval_batch_size=2,
gradient_accumulation_steps=8, num_train_epochs=3, learning_rate=5e-5, weight_decay=0.01,
warmup_steps=500, logging_steps=100, fp16=True):
"""
Set up training arguments.
"""
training_args = TrainingArguments(
output_dir=output_dir,
evaluation_strategy="epoch",
per_device_train_batch_size=per_device_train_batch_size,
per_device_eval_batch_size=per_device_eval_batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
num_train_epochs=num_train_epochs,
save_strategy="epoch",
save_total_limit=2,
logging_dir="./logs",
logging_steps=logging_steps,
report_to="none",
fp16=fp16,
learning_rate=learning_rate,
weight_decay=weight_decay,
warmup_steps=warmup_steps,
dataloader_num_workers=4,
push_to_hub=False
)
return training_args
def save_model_and_tokenizer(model, tokenizer, save_dir):
"""
Save the model and tokenizer.
"""
os.makedirs(save_dir, exist_ok=True)
model.save_pretrained(save_dir)
tokenizer.save_pretrained(save_dir)
print(f"Model and tokenizer saved at {save_dir}") |