import argparse
import os
from huggingface_hub import HfApi

TOKEN = os.environ.get("DATACOMP_TOKEN")
api = HfApi(token=TOKEN)

parser = argparse.ArgumentParser()
parser.add_argument(
    "--out_dir",
    default="teaching_arithmetic/out/",
    type=str,
    required=False,
    help="Path to the output directory.",
)
parser.add_argument(
    "--model_out_dir",
    default="teaching_arithmetic/out/addition_plain/",
    type=str,
    required=False,
    help="Path to the model output directory.",
)
args = parser.parse_args()

print("Attempting to save the Space output directory, %s" % args.out_dir)
try:
    api.upload_folder(
        folder_path=args.out_dir,
        path_in_repo=args.out_dir,
        repo_id="datacomp/teaching_arithmetic_out_directory",
        repo_type="dataset",
    )
except Exception as e:
    print("That didn't work. Error:")
    print(e)


print("Attempting to save the Space model, %s" % args.model_out_dir)
try:
    api.upload_folder(
        folder_path=args.model_out_dir,
        path_in_repo=args.model_out_dir,
        repo_id="datacomp/addition_plain",
        repo_type="model",
    )
except Exception as e:
    print("That didn't work. Error:")
    print(e)