tomaarsen HF staff commited on
Commit
278472a
1 Parent(s): 06a4764

Upload train.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train.py +38 -12
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
- model_name = "xlm-roberta-base"
 
15
  model = SpanMarkerModel.from_pretrained(
16
- model_name,
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="models/span_marker_xlm_roberta_base_fewnerd_fine_super",
 
27
  # Training Hyperparameters:
28
  learning_rate=1e-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,
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=2,
42
- dataloader_num_workers=2,
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"].select(range(3000)),
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()