tomaarsen HF staff commited on
Commit
fca338a
1 Parent(s): 4f00f10

Upload train.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train.py +60 -0
train.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datasets import load_dataset
2
+ from span_marker import SpanMarkerModel, Trainer
3
+ from transformers import TrainingArguments
4
+
5
+
6
+ def main() -> None:
7
+ # Load the dataset, ensure "tokens", "ner_tags", "document_id" and "sentence_id" columns,
8
+ # and get a list of labels
9
+ dataset = load_dataset("tomaarsen/conll2003")
10
+ labels = dataset["train"].features["ner_tags"].feature.names
11
+
12
+ # Initialize a SpanMarker model using a pretrained BERT-style encoder
13
+ model_name = "xlm-roberta-large"
14
+ model = SpanMarkerModel.from_pretrained(
15
+ model_name,
16
+ labels=labels,
17
+ # SpanMarker hyperparameters:
18
+ model_max_length=512,
19
+ marker_max_length=128,
20
+ entity_max_length=8,
21
+ )
22
+
23
+ # Prepare the 🤗 transformers training arguments
24
+ args = TrainingArguments(
25
+ output_dir="models/span_marker_xlm_roberta_large_conll03_doc_context",
26
+ # Training Hyperparameters:
27
+ learning_rate=1e-5,
28
+ per_device_train_batch_size=4,
29
+ per_device_eval_batch_size=4,
30
+ gradient_accumulation_steps=2,
31
+ num_train_epochs=3,
32
+ weight_decay=0.01,
33
+ warmup_ratio=0.1,
34
+ bf16=True,
35
+ # Other Training parameters
36
+ logging_first_step=True,
37
+ logging_steps=50,
38
+ evaluation_strategy="steps",
39
+ save_strategy="steps",
40
+ eval_steps=1000,
41
+ dataloader_num_workers=2,
42
+ )
43
+
44
+ # Initialize the trainer using our model, training args & dataset, and train
45
+ trainer = Trainer(
46
+ model=model,
47
+ args=args,
48
+ train_dataset=dataset["train"],
49
+ eval_dataset=dataset["validation"],
50
+ )
51
+ trainer.train()
52
+ trainer.save_model("models/span_marker_xlm_roberta_large_conll03_doc_context/checkpoint-final")
53
+
54
+ # Compute & save the metrics on the test set
55
+ metrics = trainer.evaluate(dataset["test"], metric_key_prefix="test")
56
+ trainer.save_metrics("test", metrics)
57
+
58
+
59
+ if __name__ == "__main__":
60
+ main()