dan_AI / train /train.py
Jack Bielinski
hi
ed390c4
raw
history blame contribute delete
689 Bytes
from transformers import TrainingArguments, Trainer, AutoModelForSequenceClassification
# Define your model
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
# Define your training arguments
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=3,
per_device_train_batch_size=16,
per_device_eval_batch_size=64,
warmup_steps=500,
weight_decay=0.01,
logging_dir='./logs',
)
# Create a Trainer instance
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset["train"],
eval_dataset=tokenized_dataset["validation"],
)
# Train the model
trainer.train()