upload data to kaggle and hugging face

Author: Dpbm

airflow-setup-test.sh

@@ -20,9 +20,11 @@ export DB_PASSWORD="airflowpass"
 export KAGGLE_USERNAME=$(clear_quotation_marks $(cat $KAGGLE_KEY_FILE | jq .username))
 export KAGGLE_KEY=$(clear_quotation_marks $(cat $KAGGLE_KEY_FILE | jq .key))
 export KAGGLE_DATASET="dpbmanalysis/quantum-circuit-images"
+export KAGGLE_MODEL="dpbmanalysis/qcop/pyTorch/standard"
 
 # HF_TOKEN comes from system
 export HF_DATASET="Dpbm/quantum-circuits"
+export HF_MODEL_REPO="Dpbm/qcop"
 
 # fix permission for volume
 sudo rm -rf ./data

compose.yml

@@ -41,8 +41,10 @@ services:
       - KAGGLE_USERNAME=${KAGGLE_USERNAME}
       - KAGGLE_KEY=${KAGGLE_KEY}
       - KAGGLE_DATASET=${KAGGLE_DATASET}
+      - KAGGLE_MODEL=${KAGGLE_MODEL}
       - HF_TOKEN=${HF_TOKEN}
       - HF_DATASET=${HF_DATASET}
+      - HF_MODEL_REPO=${HF_MODEL_REPO}
       - PYTHONPATH=/home/airflow/project
       - TARGET_FOLDER=/home/airflow/data
       - TZ=America/Sao_Paulo

dags/dataset.py

@@ -120,18 +120,16 @@
     Run training after finishing all processes.
     """
 
-    kaggle_dataset = os.environ.get("KAGGLE_DATASET")
     send_kaggle = PythonOperator(
         task_id="send_kaggle",
         python_callable=upload_dataset_kaggle,
-        op_args=[kaggle_dataset,folder]
+        op_args=[folder]
     )
 
-    hf_dataset = os.environ.get("HF_DATASET")
     send_hf = PythonOperator(
         task_id="send_huggingface",
         python_callable=upload_dataset_huggingface,
-        op_args=[hf_dataset,folder]
+        op_args=[folder]
     )
 
     send_hf.doc_md = """

dags/train.py

@@ -8,6 +8,8 @@
 from train import setup_and_run_training
 from args.parser import Arguments
 from utils.constants import DEFAULT_TARGET_FOLDER
+from export.kaggle import upload_model as upload_model_kaggle
+from export.huggingface import upload_model as upload_model_hf
 
 with DAG(dag_id="train_model", description="train vision model") as dag:
     # the env variable is meant to ease the docker image usage
@@ -16,9 +18,34 @@
     args.target_folder = folder
 
     train = PythonOperator(
-        task_id="train_model", python_callable=setup_and_run_training, op_args=[args]
+        task_id="train_model", 
+        python_callable=setup_and_run_training, 
+        op_args=[args]
     )
 
     train.doc_md = """
     Run the training cycle
     """
+
+    upload_kaggle = PythonOperator(
+        task_id="upload_kaggle",
+        python_callable=upload_model_kaggle,
+        op_args=[folder]
+    )
+
+    upload_kaggle.doc_md = """
+    Send model file to kaggle
+    """
+    
+    upload_hf = PythonOperator(
+        task_id="upload_hugginface",
+        python_callable=upload_model_hf,
+        op_args=[folder]
+    )
+
+    upload_hf.doc_md = """
+    Send model file to huggingface
+    """
+
+    train >> upload_kaggle
+    train >> upload_hf

export/helpers.py

@@ -0,0 +1,25 @@
+"""Helpers for exporting data"""
+
+import os
+
+from utils.constants import MODEL_FILE_PREFIX
+from utils.datatypes import FilePath
+
+def get_latest_model(folder:FilePath) -> FilePath:
+    """
+    Check between model files which was the latest modified (the latest model)
+    """
+
+    model_files = []
+    for file in os.listdir(folder):
+        if not file.startswith(MODEL_FILE_PREFIX):
+            continue
+        model_files.append(file)
+
+    get_file_mod_time = lambda file: os.path.getmtime(os.path.join(folder, file)) # noqa: E731
+    model_files.sort(key=get_file_mod_time, reverse=True)
+    return model_files[0]
+
+    
+
+

export/huggingface.py

@@ -1,14 +1,51 @@
 """Export dataset and model to huggingface"""
 
 import os
+import argparse
+import sys
 
 from huggingface_hub import HfApi
 
-def upload_dataset(dataset_name:str, folder:str):
+from utils.datatypes import FilePath
+from export.helpers import get_latest_model
+
+def upload_dataset(folder:FilePath):
+    """
+    Upload dataset to huggingface
+    """
+
     api = HfApi(token=os.getenv("HF_TOKEN"))
+    dataset_name = str(os.getenv("HF_DATASET"))
+
     api.upload_folder(
         folder_path=folder,
         repo_id=dataset_name,
         repo_type="dataset",
         ignore_patterns=["dataset/"]
-    )
+    )
+
+def upload_model(folder:str):
+    """
+    Get model file and upload it to huggingface
+    """
+
+    latest_model = get_latest_model(folder)
+    api = HfApi(token=os.getenv("HF_TOKEN"))
+    model_name = str(os.getenv("HF_MODEL_REPO"))
+    api.upload_file(
+        path_or_fileobj=os.path.join(folder,latest_model),
+        path_in_repo=latest_model,
+        repo_id=model_name,
+        repo_type="model",
+    )
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("path", type=str)
+    args = parser.parse_args(sys.argv[1:])
+
+    upload_dataset(args.path)
+    upload_model(args.path)
+

export/kaggle.py

@@ -1,18 +1,47 @@
 """Export dataset and model to kaggle"""
 
+import os
+import sys
 from time import ctime
+import argparse
 
 import kagglehub as kh
 
+from utils.datatypes import FilePath
 
-def upload_dataset(dataset_name:str, folder:str):
+def upload_dataset(folder:FilePath):
     """
     Upload dataset files to kaggle
     """
-
     version = ctime()
+    dataset_name = str(os.getenv("KAGGLE_DATASET"))
+
     kh.dataset_upload(
         dataset_name, 
         folder, 
         version_notes=version,
-        ignore_patterns=["dataset/"])
+        ignore_patterns=["dataset/"]
+    )
+
+def upload_model(folder:str):
+    """
+    Get model file and upload it to kaggle
+    """
+
+    version = ctime()
+    model_name = str(os.getenv("KAGGLE_MODEL"))
+
+    kh.model_upload(
+        handle=model_name,
+        local_model_dir=folder,
+        version_notes=version,
+        ignore_patterns=["dataset/","ghz*","*.zip","*.h5","*.csv"]
+    )
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("path", type=str)
+    args = parser.parse_args(sys.argv[1:])
+
+    upload_dataset(args.path)
+    upload_model(args.path)

train.py

@@ -18,6 +18,7 @@
 
 from utils.constants import (
     DEBUG,
+    MODEL_FILE_PREFIX,
     dataset_file,
     images_h5_file,
     ghz_file,
@@ -208,7 +209,7 @@ def forward(self, image: torch.Tensor) -> torch.Tensor:
 
     def save(self):
         """Save model weights."""
-        path = "model_%s" % (time.ctime())
+        path = "%s%s" % (MODEL_FILE_PREFIX, time.ctime())
         torch.save(self.state_dict(), path)
 
 

utils/constants.py

@@ -25,6 +25,8 @@
 
 DEFAULT_CHECKPOINT = None
 
+MODEL_FILE_PREFIX = "model_"
+
 # ruff: noqa: E731
 dataset_path = lambda target_folder: os.path.join(target_folder, "dataset")
 dataset_file = lambda target_folder: os.path.join(target_folder, "dataset.csv")