import tempfile from typing import List import numpy as np from huggingface_hub import HfApi, upload_folder def add_dataset_to_hub( domain: str, task: str, revision: str = "main", test_split=0.1, push_to_hub=False, text_observations: List[np.array] = None, image_observations: List[np.array] = None, discrete_observations: List[np.array] = None, continuous_observations: List[np.array] = None, discrete_actions: List[np.array] = None, continuous_actions: List[np.array] = None, rewards: List[np.array] = None, ): """ This function takes different types of observations, actions, and rewards, and prepares them for upload to the hub. It then optionally pushes them to a specified dataset repository on the data hub. Args: domain (str): The domain to which the dataset belongs. task (str): The specific task within the domain. revision (str, optional): The revision name. Default is "main". test_split (float, optional): Fraction of the dataset to be used as test data. Default is 0.1. push_to_hub (bool, optional): If True, the dataset will be pushed to the data hub. Default is False. text_observations (List[np.array], optional): List of numpy arrays with text observations. image_observations (List[np.array], optional): List of numpy arrays with image observations. discrete_observations (List[np.array], optional): List of numpy arrays with discrete observations. continuous_observations (List[np.array], optional): List of numpy arrays with continuous observations. discrete_actions (List[np.array], optional): List of numpy arrays with discrete actions. continuous_actions (List[np.array], optional): List of numpy arrays with continuous actions. rewards (List[np.array], optional): List of numpy arrays with rewards. """ n_obs = 0 n_act = 0 assert rewards is not None # assume we always have rewards assert rewards[0].dtype == np.float32, f"rewards are the of type np.float32 {rewards[0].dtype}" assert 0.0 < test_split < 1.0 n_episodes = len(rewards) train_index_end = int(n_episodes * (1.0 - test_split)) # Assume no need to shuffle the data dataset_train = {} dataset_test = {} dataset_train["rewards"] = np.array(rewards[:train_index_end], dtype=object) if test_split > 0.0: dataset_test["rewards"] = np.array(rewards[train_index_end:], dtype=object) # check the types of the arrays if text_observations is not None: assert text_observations[0].dtype == str, f"rewards are the of type np.float32 {rewards[0].dtype}" n_obs += 1 dataset_train["text_observations"] = np.array(text_observations[:train_index_end], dtype=object) if test_split > 0.0: dataset_test["text_observations"] = np.array(text_observations[train_index_end:], dtype=object) if image_observations is not None: assert ( image_observations[0].dtype == np.uint8 ), f"image_observations are the of type np.uint8 {image_observations[0].dtype}" assert len(image_observations) == n_episodes n_obs += 1 dataset_train["image_observations"] = np.array(image_observations[:train_index_end], dtype=object) if test_split > 0.0: dataset_test["image_observations"] = np.array(image_observations[train_index_end:], dtype=object) if discrete_observations is not None: assert ( discrete_observations[0].dtype == np.int64 ), f"discrete_observations are the of type np.int64 {discrete_observations[0].dtype}" assert len(discrete_observations) == n_episodes n_obs += 1 dataset_train["discrete_observations"] = np.array(discrete_observations[:train_index_end], dtype=object) if test_split > 0.0: dataset_test["discrete_observations"] = np.array(discrete_observations[train_index_end:], dtype=object) if continuous_observations is not None: assert ( continuous_observations[0].dtype == np.float32 ), f"continuous_observations are the of type np.float32 {continuous_observations[0].dtype}" assert len(continuous_observations) == n_episodes n_obs += 1 dataset_train["continuous_observations"] = np.array(continuous_observations[:train_index_end], dtype=object) if test_split > 0.0: dataset_test["continuous_observations"] = np.array(continuous_observations[train_index_end:], dtype=object) if discrete_actions is not None: assert ( discrete_actions[0].dtype == np.int64 ), f"discrete_actions are the of type np.int64 {discrete_actions[0].dtype}" assert len(discrete_actions) == n_episodes n_act += 1 dataset_train["discrete_actions"] = np.array(discrete_actions[:train_index_end], dtype=object) if test_split > 0.0: dataset_test["discrete_actions"] = np.array(discrete_actions[train_index_end:], dtype=object) if continuous_actions is not None: assert ( continuous_actions[0].dtype == np.float32 ), f"continuous_actions are the of type np.float32 {continuous_actions[0].dtype}" assert len(continuous_actions) == n_episodes n_act += 1 dataset_train["continuous_actions"] = np.array(continuous_actions[:train_index_end], dtype=object) if test_split > 0.0: dataset_test["continuous_actions"] = np.array(continuous_actions[train_index_end:], dtype=object) assert n_obs > 0, "there must be at least one observation array" assert n_act > 0, "there must be at least one action array" # TODO: loosen this contraint for text based envs? if not push_to_hub: # useful for testing other logic return with tempfile.TemporaryDirectory() as tmpdirname: np.savez_compressed(f"{tmpdirname}/train.npz", **dataset_train) if test_split > 0.0: np.savez_compressed(f"{tmpdirname}/test.npz", **dataset_test) path_in_repo = f"data/{domain}/{task}/" commit_message = f"adds {domain} {task} {n_episodes=}" HfApi().create_repo(repo_id="jat-project/jat-dataset", private=False, exist_ok=True, repo_type="dataset") upload_folder( repo_id="jat-project/jat-dataset", commit_message=commit_message, folder_path=tmpdirname, path_in_repo=path_in_repo, ignore_patterns=[".git/*"], repo_type="dataset", revision=revision, create_pr=True, ) if __name__ == "__main__": import pytest # Testing with valid inputs rewards = [np.array([1.0, 2.0, 3.0], dtype=np.float32)] continuous_observations = [np.array([0.5, 0.2, 0.8], dtype=np.float32)] discrete_actions = [np.array([0, 1, 1], dtype=np.int64)] add_dataset_to_hub( domain="domain1", task="task1", continuous_observations=continuous_observations, rewards=rewards, discrete_actions=discrete_actions, push_to_hub=False, ) # Testing with no observations with pytest.raises(AssertionError): add_dataset_to_hub( domain="domain1", task="task1", rewards=rewards, discrete_actions=discrete_actions, push_to_hub=False ) # Testing with no actions with pytest.raises(AssertionError): add_dataset_to_hub( domain="domain1", task="task1", continuous_observations=continuous_observations, rewards=rewards, push_to_hub=False, ) # Testing with invalid rewards with pytest.raises(AssertionError): add_dataset_to_hub( domain="domain1", task="task1", continuous_observations=continuous_observations, rewards=[np.array([1, 2, 3])], # rewards are not float32 discrete_actions=discrete_actions, push_to_hub=False, ) # Testing with invalid test_split with pytest.raises(AssertionError): add_dataset_to_hub( domain="domain1", task="task1", continuous_observations=continuous_observations, rewards=rewards, discrete_actions=discrete_actions, test_split=1.5, # test_split is not in range (0, 1) push_to_hub=False, )