This commit is contained in:
Jakub Pokrywka 2021-12-15 14:52:48 +01:00
parent 69cd2f5a8b
commit 8eec5016ae

View File

@ -12,6 +12,7 @@ from tqdm.auto import tqdm
BATCH_SIZE = 4
EARLY_STOPPING = 3
WARMUP_STEPS = 10_000
LR=1e-6
STEPS_EVAL = 5_000
if TEST:
@ -34,7 +35,7 @@ eval_dataloader_small = DataLoader(eval_dataset_small, batch_size=BATCH_SIZE)
eval_dataloader_full = DataLoader(eval_dataset_full, batch_size=BATCH_SIZE)
model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=1)
optimizer = AdamW(model.parameters(), lr=1e-6)
optimizer = AdamW(model.parameters(), lr=LR)
num_epochs = 5