tomaarsen HF staff commited on
Commit
2bc8e1c
1 Parent(s): feb9524

Upload train.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train.py +62 -0
train.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datasets import load_dataset
2
+ from transformers import TrainingArguments
3
+ from span_marker import SpanMarkerModel, Trainer
4
+
5
+
6
+ def main() -> None:
7
+ # Load the dataset, ensure "tokens" and "ner_tags" columns, and get a list of labels
8
+ train_dataset = load_dataset("P3ps/Cross_ner", split="train")
9
+ test_dataset = load_dataset("P3ps/Cross_ner", split="test")
10
+ labels = train_dataset.features["ner_tags"].feature.names
11
+
12
+ # Initialize a SpanMarker model using a pretrained BERT-style encoder
13
+ model_name = "bert-base-cased"
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
+ # Prepare the 🤗 transformers training arguments
24
+ args = TrainingArguments(
25
+ output_dir=f"models/span_marker_bert_base_cross_ner",
26
+ run_name=f"bb_cross_ner",
27
+ # Training Hyperparameters:
28
+ learning_rate=5e-5,
29
+ per_device_train_batch_size=32,
30
+ per_device_eval_batch_size=32,
31
+ num_train_epochs=3,
32
+ weight_decay=0.01,
33
+ warmup_ratio=0.1,
34
+ bf16=True, # Replace `bf16` with `fp16` if your hardware can't use bf16.
35
+ # Other Training parameters
36
+ logging_first_step=True,
37
+ logging_steps=50,
38
+ evaluation_strategy="steps",
39
+ save_strategy="steps",
40
+ eval_steps=200,
41
+ save_total_limit=2,
42
+ dataloader_num_workers=2,
43
+ )
44
+
45
+ # Initialize the trainer using our model, training args & dataset, and train
46
+ trainer = Trainer(
47
+ model=model,
48
+ args=args,
49
+ train_dataset=train_dataset,
50
+ eval_dataset=test_dataset,
51
+ )
52
+ trainer.train()
53
+ trainer.save_model(f"models/span_marker_bert_base_cross_ner/checkpoint-final")
54
+
55
+ # Compute & save the metrics on the test set
56
+ metrics = trainer.evaluate(test_dataset, metric_key_prefix="test")
57
+ trainer.save_metrics("test", metrics)
58
+ trainer.create_model_card()
59
+
60
+
61
+ if __name__ == "__main__":
62
+ main()