src/utils/add_artifact.py

"""
Upload file to wandb
to upload dir: 
export WANDB_PROJECT='lux-voice-processing'
wandb artifact put --type audio --name recordings_dataset data/audio_recordings 
"""

import wandb
import argparse
import os
import logging


def upload_data(args):
    """upload artifact on wandb

    Args:
        args (dict): project, name, type, local_path
    """
    if os.path.exists(args.local_path):
        run = wandb.init(
            project=args.project, job_type="add-artifact", group="add-artifact"
        )
        artifact = wandb.Artifact(name=args.name, type=args.type)
        artifact.add_file(local_path=args.local_path)
        run.log_artifact(artifact)
        run.finish()
    else:
        print(f"File does not exist: {args.local_path}")


# TODO: not working
def upload_dir(args):
    """upload dir artifact on wandb

    Args:
        args (dict): project, name, type, local_path
    """
    try:
        if os.path.isdir(args.local_path):
            run = wandb.init(
                project=args.project, job_type="add-artifact", group="add-artifact"
            )
            artifact = wandb.Artifact(name=args.name, type=args.type)
            artifact.add_dir(
                local_path="/Users/marcellopoliti/Documents/Coding/pischool/lux-voice-processing/data/audio_recordings"
            )
            run.log_artifact(artifact)
            run.finish()
        else:
            logging.error(f"Not dir: {args.local_path}")
    except Exception as e:
        logging.exception(f"Exception: {str(e)} {type(e).__name__} ")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Arguments for LLM monitoring")

    parser.add_argument(
        "--project",
        type=str,
        help="wandb project name",
        default="lux-voice-processing",
    )

    parser.add_argument(
        "--local_path",
        type=str,
        help="local path of your artifact",
        required=True,
    )

    parser.add_argument(
        "--name",
        type=str,
        help="name of your artifact",
        required=True,
    )

    parser.add_argument(
        "--type",
        type=str,
        help="type of your artifact",
        required=True,
    )

    parser.add_argument(
        "--isdir", type=bool, help="is dir?", required=False, default=False
    )
    args = parser.parse_args()

    if args.isdir:
        print("uploading dir... ", args.local_path)
        upload_data(args)
    else:
        print("uploading file...")
        upload_dir(args)