Upload train.py with huggingface_hub

train.py

@@ -1,28 +1,47 @@
+from pathlib import Path
+import shutil
 from datasets import load_dataset
 from transformers import TrainingArguments
-from span_marker import SpanMarkerModel, Trainer
+from span_marker import SpanMarkerModel, Trainer, SpanMarkerModelCardData
+
+import os
+
+os.environ["CODECARBON_LOG_LEVEL"] = "error"
 
 
 def main() -> None:
     # Load the dataset, ensure "tokens" and "ner_tags" columns, and get a list of labels
-    dataset = load_dataset("acronym_identification").rename_column("labels", "ner_tags")
+    dataset_name = "Acronym Identification"
+    dataset_id = "acronym_identification"
+    dataset = load_dataset(dataset_id).rename_column("labels", "ner_tags")
     labels = dataset["train"].features["ner_tags"].feature.names
 
     # Initialize a SpanMarker model using a pretrained BERT-style encoder
-    model_name = "bert-base-cased"
+    encoder_id = "bert-base-cased"
+    model_id = "tomaarsen/span-marker-bert-base-acronyms"
     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=dataset_name,
+            dataset_id=dataset_id,
+            license="apache-2.0",
+            language="en",
+        ),
     )
 
     # Prepare the 🤗 transformers training arguments
+    output_dir = Path("models") / model_id
     args = TrainingArguments(
-        output_dir=f"models/span_marker_bert_base_acronyms",
-        run_name=f"bb_acronyms",
+        output_dir=output_dir,
+        run_name=model_id,
         # Training Hyperparameters:
         learning_rate=5e-5,
         per_device_train_batch_size=32,
@@ -49,12 +68,13 @@ def main() -> None:
         eval_dataset=dataset["validation"],
     )
     trainer.train()
-    trainer.save_model(f"models/span_marker_bert_base_acronyms/checkpoint-final")
 
     # Compute & save the metrics on the test set
-    metrics = trainer.evaluate()
+    metrics = trainer.evaluate(metric_key_prefix="validation")
     trainer.save_metrics("validation", metrics)
-    trainer.create_model_card()
+
+    trainer.save_model(output_dir / "checkpoint-final")
+    shutil.copy2(__file__, output_dir / "checkpoint-final" / "train.py")
 
 
 if __name__ == "__main__":