Upload train.py with huggingface_hub
Browse files
train.py
CHANGED
@@ -1,6 +1,10 @@
|
|
|
|
|
|
1 |
from datasets import load_dataset
|
2 |
-
from span_marker import SpanMarkerModel, Trainer
|
3 |
from transformers import TrainingArguments
|
|
|
|
|
|
|
4 |
|
5 |
|
6 |
def main() -> None:
|
@@ -11,35 +15,46 @@ def main() -> None:
|
|
11 |
labels = dataset["train"].features["ner_tags"].feature.names
|
12 |
|
13 |
# Initialize a SpanMarker model using a pretrained BERT-style encoder
|
14 |
-
|
|
|
15 |
model = SpanMarkerModel.from_pretrained(
|
16 |
-
|
17 |
labels=labels,
|
18 |
# SpanMarker hyperparameters:
|
19 |
model_max_length=256,
|
20 |
marker_max_length=128,
|
21 |
entity_max_length=8,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
22 |
)
|
23 |
|
24 |
# Prepare the 🤗 transformers training arguments
|
|
|
25 |
args = TrainingArguments(
|
26 |
-
output_dir=
|
|
|
27 |
# Training Hyperparameters:
|
28 |
learning_rate=1e-5,
|
29 |
-
per_device_train_batch_size=
|
30 |
-
per_device_eval_batch_size=
|
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=3000,
|
41 |
-
save_total_limit=
|
42 |
-
dataloader_num_workers=
|
43 |
)
|
44 |
|
45 |
# Initialize the trainer using our model, training args & dataset, and train
|
@@ -47,15 +62,26 @@ def main() -> None:
|
|
47 |
model=model,
|
48 |
args=args,
|
49 |
train_dataset=dataset["train"],
|
50 |
-
eval_dataset=dataset["validation"]
|
51 |
)
|
52 |
trainer.train()
|
53 |
-
trainer.save_model("models/span_marker_xlm_roberta_base_fewnerd_fine_super/checkpoint-final")
|
54 |
|
55 |
# Compute & save the metrics on the test set
|
56 |
metrics = trainer.evaluate(dataset["test"], metric_key_prefix="test")
|
57 |
trainer.save_metrics("test", metrics)
|
58 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
59 |
|
60 |
if __name__ == "__main__":
|
61 |
-
main()
|
|
|
1 |
+
from pathlib import Path
|
2 |
+
import shutil
|
3 |
from datasets import load_dataset
|
|
|
4 |
from transformers import TrainingArguments
|
5 |
+
from span_marker import SpanMarkerModel, Trainer
|
6 |
+
from span_marker.model_card import SpanMarkerModelCardData
|
7 |
+
from huggingface_hub import upload_folder, upload_file
|
8 |
|
9 |
|
10 |
def main() -> None:
|
|
|
15 |
labels = dataset["train"].features["ner_tags"].feature.names
|
16 |
|
17 |
# Initialize a SpanMarker model using a pretrained BERT-style encoder
|
18 |
+
encoder_id = "xlm-roberta-base"
|
19 |
+
model_id = f"tomaarsen/span-marker-xlm-roberta-base-fewnerd-fine-super"
|
20 |
model = SpanMarkerModel.from_pretrained(
|
21 |
+
encoder_id,
|
22 |
labels=labels,
|
23 |
# SpanMarker hyperparameters:
|
24 |
model_max_length=256,
|
25 |
marker_max_length=128,
|
26 |
entity_max_length=8,
|
27 |
+
# Model card variables
|
28 |
+
model_card_data=SpanMarkerModelCardData(
|
29 |
+
model_id=model_id,
|
30 |
+
encoder_id=encoder_id,
|
31 |
+
dataset_name="FewNERD",
|
32 |
+
license="cc-by-sa-4.0",
|
33 |
+
language=["en", "multilingual"],
|
34 |
+
),
|
35 |
)
|
36 |
|
37 |
# Prepare the 🤗 transformers training arguments
|
38 |
+
output_dir = Path("models") / model_id
|
39 |
args = TrainingArguments(
|
40 |
+
output_dir=output_dir,
|
41 |
+
run_name=model_id,
|
42 |
# Training Hyperparameters:
|
43 |
learning_rate=1e-5,
|
44 |
+
per_device_train_batch_size=16,
|
45 |
+
per_device_eval_batch_size=16,
|
46 |
num_train_epochs=3,
|
47 |
weight_decay=0.01,
|
48 |
warmup_ratio=0.1,
|
49 |
+
bf16=True, # Replace `bf16` with `fp16` if your hardware can't use bf16.
|
50 |
# Other Training parameters
|
51 |
logging_first_step=True,
|
52 |
logging_steps=50,
|
53 |
evaluation_strategy="steps",
|
54 |
save_strategy="steps",
|
55 |
eval_steps=3000,
|
56 |
+
save_total_limit=1,
|
57 |
+
dataloader_num_workers=4,
|
58 |
)
|
59 |
|
60 |
# Initialize the trainer using our model, training args & dataset, and train
|
|
|
62 |
model=model,
|
63 |
args=args,
|
64 |
train_dataset=dataset["train"],
|
65 |
+
eval_dataset=dataset["validation"],
|
66 |
)
|
67 |
trainer.train()
|
|
|
68 |
|
69 |
# Compute & save the metrics on the test set
|
70 |
metrics = trainer.evaluate(dataset["test"], metric_key_prefix="test")
|
71 |
trainer.save_metrics("test", metrics)
|
72 |
|
73 |
+
# Save the model & training script locally
|
74 |
+
trainer.save_model(output_dir / "checkpoint-final")
|
75 |
+
shutil.copy2(__file__, output_dir / "checkpoint-final" / "train.py")
|
76 |
+
|
77 |
+
# Upload everything to the Hub
|
78 |
+
breakpoint()
|
79 |
+
model.push_to_hub(model_id, private=True)
|
80 |
+
upload_folder(folder_path=output_dir / "runs", path_in_repo="runs", repo_id=model_id)
|
81 |
+
upload_file(path_or_fileobj=__file__, path_in_repo="train.py", repo_id=model_id)
|
82 |
+
upload_file(path_or_fileobj=output_dir / "all_results.json", path_in_repo="all_results.json", repo_id=model_id)
|
83 |
+
upload_file(path_or_fileobj=output_dir / "emissions.csv", path_in_repo="emissions.csv", repo_id=model_id)
|
84 |
+
|
85 |
|
86 |
if __name__ == "__main__":
|
87 |
+
main()
|