nbroad HF staff commited on
Commit
ff12874
1 Parent(s): 0d773b6

Upload train.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train.py +96 -0
train.py ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ import random
3
+ import shutil
4
+ from datasets import load_dataset
5
+ from transformers import TrainingArguments
6
+ from span_marker import SpanMarkerModel, Trainer
7
+ from span_marker.model_card import SpanMarkerModelCardData
8
+ from huggingface_hub import upload_folder, upload_file
9
+
10
+
11
+
12
+ def main() -> None:
13
+ # Load the dataset, ensure "tokens" and "ner_tags" columns, and get a list of labels
14
+ labels = ["O", "B-ORG", "I-ORG"]
15
+ dataset_id = "tomaarsen/ner-orgs"
16
+ dataset = load_dataset(dataset_id)
17
+
18
+ train_dataset = dataset["train"]
19
+ eval_dataset = dataset["validation"]
20
+ eval_dataset = eval_dataset.select(random.sample(range(len(eval_dataset)), k=3000))
21
+ test_dataset = dataset["test"]
22
+
23
+ # Initialize a SpanMarker model using a pretrained BERT-style encoder
24
+ encoder_id = "BAAI/bge-base-en-v1.5"
25
+ model_id = "nbroad/span-marker-bge-base-orgs-v1"
26
+ model = SpanMarkerModel.from_pretrained(
27
+ encoder_id,
28
+ labels=labels,
29
+ # SpanMarker hyperparameters:
30
+ model_max_length=256,
31
+ marker_max_length=128,
32
+ entity_max_length=8,
33
+ # Model card variables
34
+ model_card_data=SpanMarkerModelCardData(
35
+ model_id=model_id,
36
+ dataset_id=dataset_id,
37
+ encoder_id=encoder_id,
38
+ dataset_name="FewNERD, CoNLL2003, and OntoNotes v5",
39
+ license="cc-by-sa-4.0",
40
+ language=["en"],
41
+ ),
42
+ )
43
+
44
+ # Prepare the 🤗 transformers training arguments
45
+ output_dir = Path("models") / model_id
46
+ args = TrainingArguments(
47
+ output_dir=output_dir,
48
+ run_name=model_id,
49
+ # Training Hyperparameters:
50
+ learning_rate=5e-5,
51
+ per_device_train_batch_size=128,
52
+ per_device_eval_batch_size=128,
53
+ num_train_epochs=3,
54
+ weight_decay=0.01,
55
+ warmup_ratio=0.05,
56
+ fp16=True, # Replace `bf16` with `fp16` if your hardware can't use bf16.
57
+ # Other Training parameters
58
+ logging_first_step=True,
59
+ logging_steps=100,
60
+ evaluation_strategy="steps",
61
+ save_strategy="steps",
62
+ eval_steps=600,
63
+ save_total_limit=1,
64
+ dataloader_num_workers=4,
65
+ metric_for_best_model="overall_f1",
66
+ greater_is_better=True,
67
+ )
68
+
69
+ # Initialize the trainer using our model, training args & dataset, and train
70
+ trainer = Trainer(
71
+ model=model,
72
+ args=args,
73
+ train_dataset=train_dataset,
74
+ eval_dataset=eval_dataset,
75
+ )
76
+ trainer.train()
77
+
78
+ # Compute & save the metrics on the test set
79
+ metrics = trainer.evaluate(test_dataset, metric_key_prefix="test")
80
+ trainer.save_metrics("test", metrics)
81
+
82
+ # Save the model & training script locally
83
+ trainer.save_model(output_dir / "checkpoint-final")
84
+ shutil.copy2(__file__, output_dir / "checkpoint-final" / "train.py")
85
+
86
+ # Upload everything to the Hub
87
+ breakpoint()
88
+ model.push_to_hub(model_id, private=True)
89
+ upload_folder(folder_path=output_dir / "runs", path_in_repo="runs", repo_id=model_id)
90
+ upload_file(path_or_fileobj=__file__, path_in_repo="train.py", repo_id=model_id)
91
+ upload_file(path_or_fileobj=output_dir / "all_results.json", path_in_repo="all_results.json", repo_id=model_id)
92
+ upload_file(path_or_fileobj=output_dir / "emissions.csv", path_in_repo="emissions.csv", repo_id=model_id)
93
+
94
+
95
+ if __name__ == "__main__":
96
+ main()