Upload train.py with huggingface_hub

train.py

@@ -1,6 +1,10 @@
+from pathlib import Path
+import shutil
 from datasets import load_dataset
-from span_marker import SpanMarkerModel, Trainer
 from transformers import TrainingArguments
+from span_marker import SpanMarkerModel, Trainer
+from span_marker.model_card import SpanMarkerModelCardData
+from huggingface_hub import upload_folder, upload_file
 
 
 def main() -> None:
@@ -11,35 +15,46 @@ def main() -> None:
     labels = dataset["train"].features["ner_tags"].feature.names
 
     # Initialize a SpanMarker model using a pretrained BERT-style encoder
-    model_name = "xlm-roberta-base"
+    encoder_id = "xlm-roberta-base"
+    model_id = f"tomaarsen/span-marker-xlm-roberta-base-fewnerd-fine-super"
     model = SpanMarkerModel.from_pretrained(
-        model_name,
+        encoder_id,
         labels=labels,
         # SpanMarker hyperparameters:
         model_max_length=256,
         marker_max_length=128,
         entity_max_length=8,
+        # Model card variables
+        model_card_data=SpanMarkerModelCardData(
+            model_id=model_id,
+            encoder_id=encoder_id,
+            dataset_name="FewNERD",
+            license="cc-by-sa-4.0",
+            language=["en", "multilingual"],
+        ),
     )
 
     # Prepare the 🤗 transformers training arguments
+    output_dir = Path("models") / model_id
     args = TrainingArguments(
-        output_dir="models/span_marker_xlm_roberta_base_fewnerd_fine_super",
+        output_dir=output_dir,
+        run_name=model_id,
         # Training Hyperparameters:
         learning_rate=1e-5,
-        per_device_train_batch_size=32,
-        per_device_eval_batch_size=32,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
         num_train_epochs=3,
         weight_decay=0.01,
         warmup_ratio=0.1,
-        bf16=True,
+        bf16=True,  # Replace `bf16` with `fp16` if your hardware can't use bf16.
         # Other Training parameters
         logging_first_step=True,
         logging_steps=50,
         evaluation_strategy="steps",
         save_strategy="steps",
         eval_steps=3000,
-        save_total_limit=2,
-        dataloader_num_workers=2,
+        save_total_limit=1,
+        dataloader_num_workers=4,
     )
 
     # Initialize the trainer using our model, training args & dataset, and train
@@ -47,15 +62,26 @@ def main() -> None:
         model=model,
         args=args,
         train_dataset=dataset["train"],
-        eval_dataset=dataset["validation"].select(range(3000)),
+        eval_dataset=dataset["validation"],
     )
     trainer.train()
-    trainer.save_model("models/span_marker_xlm_roberta_base_fewnerd_fine_super/checkpoint-final")
 
     # Compute & save the metrics on the test set
     metrics = trainer.evaluate(dataset["test"], metric_key_prefix="test")
     trainer.save_metrics("test", metrics)
 
+    # Save the model & training script locally
+    trainer.save_model(output_dir / "checkpoint-final")
+    shutil.copy2(__file__, output_dir / "checkpoint-final" / "train.py")
+
+    # Upload everything to the Hub
+    breakpoint()
+    model.push_to_hub(model_id, private=True)
+    upload_folder(folder_path=output_dir / "runs", path_in_repo="runs", repo_id=model_id)
+    upload_file(path_or_fileobj=__file__, path_in_repo="train.py", repo_id=model_id)
+    upload_file(path_or_fileobj=output_dir / "all_results.json", path_in_repo="all_results.json", repo_id=model_id)
+    upload_file(path_or_fileobj=output_dir / "emissions.csv", path_in_repo="emissions.csv", repo_id=model_id)
+
 
 if __name__ == "__main__":
-    main()
+    main()