fix: [training] Push CodeCarbon emissions report to the Hub

The emissions.csv was written to the CWD (CodeCarbon's default output_dir),
which trainer.push_to_hub() never uploads, so the report stopped reaching the
Hub. Point the tracker at model_save_dir and push the report explicitly via a
shared push_emissions_report() helper.

Also convert generation_description.py and cwe_guesser_patches.py from the
@track_emissions decorator to an explicit EmissionsTracker: the decorator only
finalizes emissions.csv after train() returns, i.e. after push_to_hub() has
already run, so the report could never be included.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: cedricbonhomme

vulntrain/trainers/classify_severity.py

@@ -16,6 +16,8 @@
     TrainingArguments,
 )
 
+from vulntrain.utils import push_emissions_report
+
 """
 Automatically classify new vulnerabilities based on their descriptions,
 even if they don't have CVSS scores.
@@ -225,7 +227,15 @@ def tokenize_function(elem):
     # Train model
     tracker = None
     if codecarbon:
-        tracker = EmissionsTracker(project_name="VulnTrain", allow_multiple_runs=True)
+        # Save emissions data inside the model directory so it gets pushed to
+        # the Hub together with the model (default output_dir is the CWD, which
+        # is never uploaded).
+        tracker = EmissionsTracker(
+            project_name="VulnTrain",
+            output_dir=model_save_dir,
+            output_file="emissions.csv",
+            allow_multiple_runs=True,
+        )
         tracker.start()
     try:
         trainer.train()
@@ -239,6 +249,9 @@ def tokenize_function(elem):
         trainer.push_to_hub()
         tokenizer.push_to_hub(repo_id)
 
+        if tracker is not None and push_emissions_report(model_save_dir, repo_id):
+            logger.info(f"Emissions report pushed to {repo_id}")
+
 
 def main():
     parser = argparse.ArgumentParser(

vulntrain/trainers/classify_severity_cnvd.py

@@ -19,6 +19,8 @@
     TrainingArguments,
 )
 
+from vulntrain.utils import push_emissions_report
+
 # Logging setup
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -311,7 +313,15 @@ def tokenize_function(examples):
     else:
         trainer = Trainer(**trainer_kwargs)
 
-    tracker = EmissionsTracker(project_name="VulnTrain", allow_multiple_runs=True)
+    # Save emissions data inside the model directory so it gets pushed to the
+    # Hub together with the model (default output_dir is the CWD, which is never
+    # uploaded).
+    tracker = EmissionsTracker(
+        project_name="VulnTrain",
+        output_dir=model_save_dir,
+        output_file="emissions.csv",
+        allow_multiple_runs=True,
+    )
     tracker.start()
     try:
         trainer.train()
@@ -323,6 +333,9 @@ def tokenize_function(examples):
     trainer.push_to_hub()
     tokenizer.push_to_hub(repo_id)
 
+    if push_emissions_report(model_save_dir, repo_id):
+        logger.info(f"Emissions report pushed to {repo_id}")
+
     if not push_card:
         return
 

vulntrain/trainers/cwe_guesser_patches.py

@@ -10,7 +10,7 @@
 import evaluate
 import numpy as np
 import torch
-from codecarbon import track_emissions
+from codecarbon import EmissionsTracker
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score, f1_score
 from transformers import (
@@ -21,6 +21,8 @@
     TrainingArguments,
 )
 
+from vulntrain.utils import push_emissions_report
+
 accuracy = evaluate.load("accuracy")
 f1 = evaluate.load("f1", config_name="macro")
 
@@ -62,7 +64,6 @@ def compute_loss(
         return (loss, outputs) if return_outputs else loss
 
 
-@track_emissions(project_name="VulnTrain", allow_multiple_runs=True)
 def train(base_model, dataset_id, repo_id, model_save_dir="./vulnerability-classify"):
     dataset = load_dataset(dataset_id)
     dataset = dataset["train"].filter(lambda x: x.get("cwe") and len(x["cwe"]) > 0)
@@ -179,9 +180,20 @@ def tokenize_function(examples):
         class_weights=class_weights,
     )
 
+    # Save emissions data inside the model directory so it gets pushed to the
+    # Hub together with the model (default output_dir is the CWD, which is never
+    # uploaded).
+    tracker = EmissionsTracker(
+        project_name="VulnTrain",
+        output_dir=model_save_dir,
+        output_file="emissions.csv",
+        allow_multiple_runs=True,
+    )
+    tracker.start()
     try:
         trainer.train()
     finally:
+        tracker.stop()
         model.save_pretrained(model_save_dir)
         tokenizer.save_pretrained(model_save_dir)
 
@@ -200,6 +212,9 @@ def tokenize_function(examples):
     trainer.push_to_hub(repo_id)
     tokenizer.push_to_hub(repo_id)
 
+    if push_emissions_report(model_save_dir, repo_id):
+        logger.info(f"Emissions report pushed to {repo_id}")
+
 
 def main():
     parser = argparse.ArgumentParser(

vulntrain/trainers/generation_description.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 import torch
-from codecarbon import track_emissions
+from codecarbon import EmissionsTracker
 from datasets import load_dataset
 from transformers import (
     AutoModelForCausalLM,
@@ -14,6 +14,8 @@
     TrainingArguments,
 )
 
+from vulntrain.utils import push_emissions_report
+
 """
 Create a text generation model for descriptions of vulnerabilities.
 
@@ -38,7 +40,6 @@ def tokenize_function(examples):
     return tokenized_datasets.train_test_split(test_size=0.2)
 
 
-@track_emissions(project_name="VulnTrain", allow_multiple_runs=True)
 def train(
     base_model, dataset_id, repo_id, model_save_dir="./vulnerability-description"
 ):
@@ -89,15 +90,29 @@ def train(
         data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
     )
 
+    # Save emissions data inside the model directory so it gets pushed to the
+    # Hub together with the model (default output_dir is the CWD, which is never
+    # uploaded).
+    tracker = EmissionsTracker(
+        project_name="VulnTrain",
+        output_dir=model_save_dir,
+        output_file="emissions.csv",
+        allow_multiple_runs=True,
+    )
+    tracker.start()
     try:
         trainer.train()
     finally:
+        tracker.stop()
         model.save_pretrained(model_save_dir)
         tokenizer.save_pretrained(model_save_dir)
 
     trainer.push_to_hub()
     tokenizer.push_to_hub(repo_id)
 
+    if push_emissions_report(model_save_dir, repo_id):
+        print(f"Emissions report pushed to {repo_id}")
+
 
 def main():
     parser = argparse.ArgumentParser(

vulntrain/utils.py

@@ -2,7 +2,6 @@
 import re
 
 import cvss
-
 from markdown_it import MarkdownIt
 from nltk.tokenize import sent_tokenize
 
@@ -206,3 +205,28 @@ def extract_cpe(branches):
 
     # Print unique CPEs
     return sorted(set(cpe_list))
+
+
+def push_emissions_report(model_save_dir: str, repo_id: str) -> bool:
+    """Upload the CodeCarbon ``emissions.csv`` report to the Hugging Face Hub.
+
+    ``trainer.push_to_hub()`` uploads ``model_save_dir`` but ignores files it
+    does not recognise, so the emissions report is pushed explicitly.
+
+    Returns ``True`` if the report existed and was uploaded.
+    """
+    from pathlib import Path
+
+    emissions_file = Path(model_save_dir) / "emissions.csv"
+    if not emissions_file.exists():
+        return False
+
+    from huggingface_hub import HfApi
+
+    HfApi().upload_file(
+        path_or_fileobj=str(emissions_file),
+        path_in_repo="emissions.csv",
+        repo_id=repo_id,
+        commit_message="Add CodeCarbon emissions report",
+    )
+    return True