import pickle


def load_dataset(args):
    """load a pickle file and reduce samples"""
    # load data
    with open(args.input_file, "rb") as f_in:
        data = pickle.load(f_in)

    # use less data to safe time for testing
    if args.small is not None:
        # if limit is given
        max_length = len(data["features"])
        limit = min(args.small, max_length)
        # go through data and limit it
        for key, value in data.items():
            data[key] = value[:limit]

    return data