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")