from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset
Загрузка токенизатора и модели
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
Загрузка и подготовка данных
dataset = load_dataset("imdb") # Вы можете использовать свой набор данных
encoded_dataset = dataset.map(lambda examples: tokenizer(examples['text'], padding="max_length", truncation=True), batched=True)
Настройка параметров тренировки
training_args = TrainingArguments(
output_dir="./results",
evaluation_strategy="epoch",
learning_rate=2e-5,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
num_train_epochs=3,
weight_decay=0.01,
)
Создание тренера
trainer = Trainer(
model=model,
args=training_args,
train_dataset=encoded_dataset['train'],
eval_dataset=encoded_dataset['test'],
)
Тренировка модели
trainer.train()
Сохранение модели
model.save_pretrained("./my_custom_model")
tokenizer.save_pretrained("./my_custom_model")
