romainf/

distilbert-base-uncased-imdb-3000

training_args = TrainingArguments(
    output_dir="bert_results_imdb",
    learning_rate=1e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    weight_decay=0.01,
    warmup_ratio = 0.06,
    max_steps = 5000,
    optim = 'adamw_torch',
    save_strategy = 'steps',
    evaluation_strategy='steps',
    load_best_model_at_end=True
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_imdb["train"],
    eval_dataset=tokenized_imdb["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)