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

utils: add evaluator script

Browse files
Files changed (1) hide show
  1. evaluator.py +43 -0
evaluator.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sys
2
+
3
+ from datasets import load_dataset
4
+ from transformers import TrainingArguments
5
+ from span_marker import SpanMarkerModel, Trainer
6
+
7
+
8
+ # Load the dataset, ensure "tokens" and "ner_tags" columns, and get a list of labels
9
+ dataset = load_dataset("gwlms/germeval2014")
10
+ labels = dataset["train"].features["ner_tags"].feature.names
11
+
12
+ # Initialize a SpanMarker model using a pretrained BERT-style encoder
13
+ model_name = sys.argv[1]
14
+ model = SpanMarkerModel.from_pretrained(
15
+ model_name,
16
+ labels=labels,
17
+ # SpanMarker hyperparameters:
18
+ model_max_length=256,
19
+ marker_max_length=128,
20
+ entity_max_length=8,
21
+ )
22
+
23
+ args = TrainingArguments(
24
+ output_dir="/tmp",
25
+ per_device_eval_batch_size=64,
26
+ )
27
+
28
+ # Initialize the trainer using our model, training args & dataset, and train
29
+ trainer = Trainer(
30
+ model=model,
31
+ args=args,
32
+ train_dataset=dataset["train"],
33
+ eval_dataset=dataset["validation"],
34
+ )
35
+
36
+
37
+ print("Evaluating on development set...")
38
+ dev_metrics = trainer.evaluate(dataset["validation"], metric_key_prefix="eval")
39
+ print(dev_metrics)
40
+
41
+ print("Evaluating on test set...")
42
+ test_metrics = trainer.evaluate(dataset["test"], metric_key_prefix="test")
43
+ print(test_metrics)