stefan-it commited on
Commit
d06a3fa
1 Parent(s): a2af8d0

utils: add trainer script

Browse files
Files changed (1) hide show
  1. trainer.py +61 -0
trainer.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datasets import load_dataset
2
+ from transformers import TrainingArguments
3
+ from span_marker import SpanMarkerModel, Trainer
4
+
5
+ def perform_training(learning_rate: float, seed: int) -> None:
6
+ # Load the dataset, ensure "tokens" and "ner_tags" columns, and get a list of labels
7
+ dataset = load_dataset("gwlms/germeval2014")
8
+ labels = dataset["train"].features["ner_tags"].feature.names
9
+
10
+ # Initialize a SpanMarker model using a pretrained BERT-style encoder
11
+ model_name = "gwlms/span-marker-teams-germeval14"
12
+ model = SpanMarkerModel.from_pretrained(
13
+ model_name,
14
+ labels=labels,
15
+ # SpanMarker hyperparameters:
16
+ model_max_length=256,
17
+ marker_max_length=128,
18
+ entity_max_length=8,
19
+ )
20
+
21
+ # Prepare the 🤗 transformers training arguments
22
+ args = TrainingArguments(
23
+ output_dir=f"./span_marker-{model_name}-bs16-lr{learning_rate}-{seed}",
24
+ # Training Hyperparameters:
25
+ learning_rate=learning_rate,
26
+ per_device_train_batch_size=16,
27
+ per_device_eval_batch_size=16,
28
+ num_train_epochs=3,
29
+ weight_decay=0.01,
30
+ warmup_ratio=0.1,
31
+ fp16=True, # Replace `bf16` with `fp16` if your hardware can't use bf16.
32
+ # Other Training parameters
33
+ logging_first_step=True,
34
+ logging_steps=50,
35
+ evaluation_strategy="epoch",
36
+ save_strategy="epoch",
37
+ save_total_limit=11,
38
+ dataloader_num_workers=2,
39
+ seed=seed,
40
+ load_best_model_at_end=True,
41
+ )
42
+
43
+ # Initialize the trainer using our model, training args & dataset, and train
44
+ trainer = Trainer(
45
+ model=model,
46
+ args=args,
47
+ train_dataset=dataset["train"],
48
+ eval_dataset=dataset["validation"],
49
+ )
50
+ trainer.train()
51
+ trainer.save_model(f"./span_marker-{model_name}-bs16-lr{learning_rate}-{seed}/best-checkpoint")
52
+
53
+ # Compute & save the metrics on the test set
54
+ metrics = trainer.evaluate(dataset["test"], metric_key_prefix="test")
55
+ trainer.save_metrics("test", metrics)
56
+
57
+
58
+ if __name__ == "__main__":
59
+ for learning_rate in [5e-05]:
60
+ for seed in [1,2,3,4,5]:
61
+ perform_training(learning_rate, seed)