Ideas for Improvement

Repeated code snippets

Dataset IO

Local/Remote loading option (with test option)

option = 'local'

if option == 'local':
    with open(local ,'rb') as file:
        bg = pickle.load(file)
elif option == 'remote':
    with fs.open(remote, 'rb') as file:
        bg = pickle.load(file)

And a separate test option if it's a big file where test_file is a smaller version of the dataset:

option = 'local'

if option == 'local':
    with open(local_file ,'rb') as file:
        bg = pickle.load(file)
elif option == 'remote':
    with fs.open(remote_file, 'rb') as file:
        bg = pickle.load(file)
elif option == 'test':
    with fs.open(test_file, 'rb') as file:
        bg = pickle.load(file)
    if not fs.exists(remote_file):
        raise Exception(f"file {remote_file} does not exist")
    if not local_file.exists():
        raise Exception(f"file {local_file} does not exist")

Split apart and simplify

Use layers.

• Environment
• Local Config
• Notebooks
• Datasets