diff --git a/learnware/client/learnware_client.py b/learnware/client/learnware_client.py index 56c9e1d..f16cfa0 100644 --- a/learnware/client/learnware_client.py +++ b/learnware/client/learnware_client.py @@ -67,6 +67,7 @@ class LearnwareClient: self.chunk_size = 1024 * 1024 self.tempdir_list = [] + self.login_status = False atexit.register(self.cleanup) def login(self, email, token): @@ -80,7 +81,11 @@ class LearnwareClient: token = result["data"]["token"] self.headers = {"Authorization": f"Bearer {token}"} - + self.login_status = True + + def is_login(self): + return self.login_status + @require_login def logout(self): url = f"{self.host}/auth/logout" diff --git a/learnware/specification/module.py b/learnware/specification/module.py index 8938ff4..10b4fdc 100644 --- a/learnware/specification/module.py +++ b/learnware/specification/module.py @@ -175,7 +175,7 @@ def generate_rkme_text_spec( def generate_stat_spec( type: str, X: Union[np.ndarray, pd.DataFrame, torch.Tensor, List[str]], *args, **kwargs -) -> BaseStatSpecification: +) -> Union[RKMETableSpecification, RKMEImageSpecification, RKMETextSpecification]: """ Interface for users to generate statistical specification. Return a StatSpecification object, use .save() method to save as npy file. diff --git a/learnware/tests/__init__.py b/learnware/tests/__init__.py index a048b3f..5019465 100644 --- a/learnware/tests/__init__.py +++ b/learnware/tests/__init__.py @@ -1 +1,2 @@ from .module import get_semantic_specification +from .utils import parametrize \ No newline at end of file diff --git a/learnware/tests/utils.py b/learnware/tests/utils.py new file mode 100644 index 0000000..d950bf3 --- /dev/null +++ b/learnware/tests/utils.py @@ -0,0 +1,9 @@ +import unittest + +def parametrize(test_class, **kwargs): + test_loader = unittest.TestLoader() + test_names = test_loader.getTestCaseNames(test_class) + _suite = unittest.TestSuite() + for name in test_names: + _suite.addTest(test_class(name, **kwargs)) + return _suite \ No newline at end of file diff --git a/tests/test_learnware_client/test_all_learnware.py b/tests/test_learnware_client/test_all_learnware.py index 8bc9dab..f7fc9d8 100644 --- a/tests/test_learnware_client/test_all_learnware.py +++ b/tests/test_learnware_client/test_all_learnware.py @@ -3,32 +3,27 @@ import json import zipfile import unittest import tempfile +import argparse from learnware.client import LearnwareClient from learnware.specification import generate_semantic_spec from learnware.market import BaseUserInfo - +from learnware.tests import parametrize class TestAllLearnware(unittest.TestCase): - def setUp(self): - unittest.TestCase.setUpClass() - dir_path = os.path.dirname(__file__) - config_path = os.path.join(dir_path, "config.json") - if not os.path.exists(config_path): - data = {"email": None, "token": None} - with open(config_path, "w") as file: - json.dump(data, file) - - with open(config_path, "r") as file: - data = json.load(file) - email = data["email"] - token = data["token"] - - if email is None or token is None: - raise ValueError("Please set email and token in config.json.") - self.client = LearnwareClient() - self.client.login(email, token) - + client = LearnwareClient() + + def __init__(self, method_name='runTest', email=None, token=None): + super(TestAllLearnware, self).__init__(method_name) + self.email = email + self.token = token + + if self.email is not None and self.token is not None: + self.client.login(self.email, self.token) + else: + print("Client doest not login, all tests will be ignored!") + + @unittest.skipIf(not client.is_login(), "Client doest not login!") def test_all_learnware(self): max_learnware_num = 1000 semantic_spec = generate_semantic_spec() @@ -57,4 +52,10 @@ class TestAllLearnware(unittest.TestCase): if __name__ == "__main__": - unittest.main() + parser = argparse.ArgumentParser() + parser.add_argument("--email", type=str, required=False, help="The email to login learnware client") + parser.add_argument("--token", type=str, required=False, help="The token to login learnware client") + args = parser.parse_args() + + runner = unittest.TextTestRunner() + runner.run(parametrize(TestAllLearnware, email=args.email, token=args.token)) \ No newline at end of file diff --git a/tests/test_learnware_client/test_container.py b/tests/test_learnware_client/test_container.py new file mode 100644 index 0000000..c96d2ab --- /dev/null +++ b/tests/test_learnware_client/test_container.py @@ -0,0 +1,54 @@ +import os +import unittest +import argparse +import numpy as np + +from learnware.learnware import get_learnware_from_dirpath +from learnware.client import LearnwareClient +from learnware.client.container import ModelCondaContainer, LearnwaresContainer +from learnware.tests import parametrize + +class TestContainer(unittest.TestCase): + def __init__(self, method_name='runTest', mode="all"): + super(TestContainer, self).__init__(method_name) + self.modes = [] + if mode in {"all", "conda"}: + self.modes.append("conda") + if mode in {"all", "docker"}: + self.modes.append("docker") + + def setUp(self): + self.client = LearnwareClient() + + def _test_container_with_pip(self, mode): + learnware_id = "00000147" + learnware = self.client.load_learnware(learnware_id=learnware_id) + with LearnwaresContainer(learnware, ignore_error=False, mode=mode) as env_container: + learnware = env_container.get_learnwares_with_container()[0] + input_array = np.random.random(size=(20, 23)) + print(learnware.predict(input_array)) + + def _test_container_with_conda(self, mode): + learnware_id = "00000148" + learnware = self.client.load_learnware(learnware_id=learnware_id) + with LearnwaresContainer(learnware, ignore_error=False, mode=mode) as env_container: + learnware = env_container.get_learnwares_with_container()[0] + input_array = np.random.random(size=(20, 204)) + print(learnware.predict(input_array)) + + def test_container_with_pip(self): + for mode in self.modes: + self._test_container_with_pip(mode=mode) + + def test_container_with_conda(self): + for mode in self.modes: + self._test_container_with_conda(mode=mode) + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--mode", type=str, required=False, default="all", help="The mode to run container, must be in ['all', 'conda', 'docker']") + args = parser.parse_args() + + assert args.mode in {"all", "conda", "docker"}, f"The mode must be in ['all', 'conda', 'docker'], instead of '{args.mode}'" + runner = unittest.TextTestRunner() + runner.run(parametrize(TestContainer, mode=args.mode)) \ No newline at end of file diff --git a/tests/test_learnware_client/test_load_learnware.py b/tests/test_learnware_client/test_load_learnware.py index 8381f90..fafb261 100644 --- a/tests/test_learnware_client/test_load_learnware.py +++ b/tests/test_learnware_client/test_load_learnware.py @@ -8,94 +8,59 @@ from learnware.learnware import get_learnware_from_dirpath from learnware.client import LearnwareClient from learnware.client.container import ModelCondaContainer, LearnwaresContainer from learnware.reuse import AveragingReuser +from learnware.tests import parametrize -class TestLearnwareLoadWithConda(unittest.TestCase): +class TestLearnwareLoad(unittest.TestCase): + def __init__(self, method_name='runTest', mode="conda"): + super(TestLearnwareLoad, self).__init__(method_name) + self.runnable_options = [] + if mode in {"all", "conda"}: + self.runnable_options.append("conda") + if mode in {"all", "docker"}: + self.runnable_options.append("docker") + def setUp(self): self.client = LearnwareClient() root = os.path.dirname(__file__) self.learnware_ids = ["00000084", "00000154", "00000155"] self.zip_paths = [os.path.join(root, x) for x in ["1.zip", "2.zip", "3.zip"]] - self.runnable_option = "conda" - #def test_load_single_learnware_by_zippath(self): - # for learnware_id, zip_path in zip(self.learnware_ids, self.zip_paths): - # self.client.download_learnware(learnware_id, zip_path) -# - # learnware_list = [ - # self.client.load_learnware(learnware_path=zippath, runnable_option=self.runnable_option) for zippath in self.zip_paths - # ] - # reuser = AveragingReuser(learnware_list, mode="vote_by_label") - # input_array = np.random.random(size=(20, 13)) - # print(reuser.predict(input_array)) -# - # for learnware in learnware_list: - # print(learnware.id, learnware.predict(input_array)) -# - #def test_load_multi_learnware_by_zippath(self): - # for learnware_id, zip_path in zip(self.learnware_ids, self.zip_paths): - # self.client.download_learnware(learnware_id, zip_path) -# - # learnware_list = self.client.load_learnware(learnware_path=self.zip_paths, runnable_option=self.runnable_option) - # reuser = AveragingReuser(learnware_list, mode="vote_by_label") - # input_array = np.random.random(size=(20, 13)) - # print(reuser.predict(input_array)) -# - # for learnware in learnware_list: - # print(learnware.id, learnware.predict(input_array)) -# - #def test_load_single_learnware_by_id(self): - # learnware_list = [ - # self.client.load_learnware(learnware_id=idx, runnable_option=self.runnable_option) for idx in self.learnware_ids - # ] - # reuser = AveragingReuser(learnware_list, mode="vote_by_label") - # input_array = np.random.random(size=(20, 13)) - # print(reuser.predict(input_array)) -# - # for learnware in learnware_list: - # print(learnware.id, learnware.predict(input_array)) -# - #def test_load_multi_learnware_by_id(self): - # learnware_list = self.client.load_learnware(learnware_id=self.learnware_ids, runnable_option=self.runnable_option) - # reuser = AveragingReuser(learnware_list, mode="vote_by_label") - # input_array = np.random.random(size=(20, 13)) - # print(reuser.predict(input_array)) -# - # for learnware in learnware_list: - # print(learnware.id, learnware.predict(input_array)) -# - def test_load_single_learnware_by_id_pip(self): - learnware_id = "00000147" - learnware = self.client.load_learnware(learnware_id=learnware_id, runnable_option=self.runnable_option) - input_array = np.random.random(size=(20, 23)) - print(learnware.predict(input_array)) -# - def test_load_single_learnware_by_id_conda(self): - learnware_id = "00000148" - learnware = self.client.load_learnware(learnware_id=learnware_id, runnable_option=self.runnable_option) - input_array = np.random.random(size=(20, 204)) - print(learnware.predict(input_array)) -# -# -class TestLearnwareLoadWithDocker(TestLearnwareLoadWithConda): - def setUp(self): - super(TestLearnwareLoadWithDocker, self).setUp() - self.runnable_option = "docker" + def _test_load_learnware_by_zippath(self, runnable_option): + for learnware_id, zip_path in zip(self.learnware_ids, self.zip_paths): + self.client.download_learnware(learnware_id, zip_path) + + learnware_list = self.client.load_learnware(learnware_path=self.zip_paths, runnable_option=runnable_option) + reuser = AveragingReuser(learnware_list, mode="vote_by_label") + input_array = np.random.random(size=(20, 13)) + print(reuser.predict(input_array)) + for learnware in learnware_list: + print(learnware.id, learnware.predict(input_array)) + + + def _test_load_learnware_by_id(self, runnable_option): + learnware_list = self.client.load_learnware(learnware_id=self.learnware_ids, runnable_option=runnable_option) + reuser = AveragingReuser(learnware_list, mode="vote_by_label") + input_array = np.random.random(size=(20, 13)) + print(reuser.predict(input_array)) + + for learnware in learnware_list: + print(learnware.id, learnware.predict(input_array)) -def suite(mode): - _suite = unittest.TestSuite() - #_suite.addTest(TestLearnwareLoadWithDocker()) - if mode == "all" or mode == "conda": - _suite.addTest(unittest.makeSuite(TestLearnwareLoadWithConda)) - if mode == "all" or mode == "docker": - _suite.addTest(unittest.makeSuite(TestLearnwareLoadWithDocker)) - return _suite + def test_load_learnware_by_zippath(self): + for runnable_option in self.runnable_options: + self._test_load_learnware_by_zippath(runnable_option=runnable_option) + + def test_load_learnware_by_id(self): + for runnable_option in self.runnable_options: + self._test_load_learnware_by_id(runnable_option=runnable_option) + if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--mode", type=str, required=False, default="all", help="The mode to run load learnware, must be in ['all', 'conda', 'docker']") + parser.add_argument("--mode", type=str, required=False, default="conda", help="The mode to load learnware, must be in ['all', 'conda', 'docker']") args = parser.parse_args() assert args.mode in {"all", "conda", "docker"}, f"The mode must be in ['all', 'conda', 'docker'], instead of '{args.mode}'" runner = unittest.TextTestRunner() - runner.run(suite(args.mode)) \ No newline at end of file + runner.run(parametrize(TestLearnwareLoad, mode=args.mode)) \ No newline at end of file diff --git a/tests/test_learnware_client/test_upload.py b/tests/test_learnware_client/test_upload.py index 324ea77..8bcd988 100644 --- a/tests/test_learnware_client/test_upload.py +++ b/tests/test_learnware_client/test_upload.py @@ -1,32 +1,26 @@ import os -import json +import argparse import unittest import tempfile from learnware.client import LearnwareClient from learnware.specification import generate_semantic_spec +from learnware.tests import parametrize +class TestUpload(unittest.TestCase): + client = LearnwareClient() + + def __init__(self, method_name='runTest', email=None, token=None): + super(TestUpload, self).__init__(method_name) + self.email = email + self.token = token + + if self.email is not None and self.token is not None: + self.client.login(self.email, self.token) + else: + print("Client doest not login, all tests will be ignored!") -class TestAllLearnware(unittest.TestCase): - def setUp(self): - unittest.TestCase.setUpClass() - dir_path = os.path.dirname(__file__) - config_path = os.path.join(dir_path, "config.json") - if not os.path.exists(config_path): - data = {"email": None, "token": None} - with open(config_path, "w") as file: - json.dump(data, file) - - with open(config_path, "r") as file: - data = json.load(file) - email = data["email"] - token = data["token"] - - if email is None or token is None: - raise ValueError("Please set email and token in config.json.") - self.client = LearnwareClient() - self.client.login(email, token) - + @unittest.skipIf(not client.is_login(), "Client doest not login!") def test_upload(self): input_description = { "Dimension": 13, @@ -67,4 +61,10 @@ class TestAllLearnware(unittest.TestCase): if __name__ == "__main__": - unittest.main() + parser = argparse.ArgumentParser() + parser.add_argument("--email", type=str, required=False, help="The email to login learnware client") + parser.add_argument("--token", type=str, required=False, help="The token to login learnware client") + args = parser.parse_args() + + runner = unittest.TextTestRunner() + runner.run(parametrize(TestUpload, email=args.email, token=args.token)) \ No newline at end of file diff --git a/tests/test_specification/test_hetero_spec.py b/tests/test_specification/test_hetero_spec.py new file mode 100644 index 0000000..21563b3 --- /dev/null +++ b/tests/test_specification/test_hetero_spec.py @@ -0,0 +1,43 @@ +import os +import json +import string +import random +import torch +import unittest +import tempfile +import numpy as np + +from learnware.specification import RKMETableSpecification, HeteroMapTableSpecification +from learnware.specification import generate_stat_spec +from learnware.market.heterogeneous.organizer import HeteroMap + +class TestTableRKME(unittest.TestCase): + + def setUp(self): + self.hetero_map = HeteroMap() + + def _test_hetero_spec(self, X): + rkme: RKMETableSpecification = generate_stat_spec(type="table", X=X) + hetero_spec = self.hetero_map.hetero_mapping(rkme_spec=rkme, features=dict()) + with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir: + rkme_path = os.path.join(tempdir, "rkme.json") + hetero_spec.save(rkme_path) + + with open(rkme_path, "r") as f: + data = json.load(f) + assert data["type"] == "HeteroMapTableSpecification" + + rkme2 = HeteroMapTableSpecification() + rkme2.load(rkme_path) + assert rkme2.type == "HeteroMapTableSpecification" + + + def test_hetero_rkme(self): + self._test_hetero_spec(np.random.uniform(-10000, 10000, size=(5000, 200))) + self._test_hetero_spec(np.random.uniform(-10000, 10000, size=(10000, 100))) + self._test_hetero_spec(np.random.uniform(-10000, 10000, size=(5, 20))) + self._test_hetero_spec(np.random.uniform(-10000, 10000, size=(1, 50))) + self._test_hetero_spec(np.random.uniform(-10000, 10000, size=(100, 150))) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_specification/test_image_rkme.py b/tests/test_specification/test_image_rkme.py new file mode 100644 index 0000000..ad5b1ac --- /dev/null +++ b/tests/test_specification/test_image_rkme.py @@ -0,0 +1,40 @@ +import os +import json +import string +import random +import torch +import unittest +import tempfile +import numpy as np + +from learnware.specification import RKMEImageSpecification +from learnware.specification import generate_stat_spec + + +class TestImageRKME(unittest.TestCase): + @staticmethod + def _test_image_rkme(X): + image_rkme = generate_stat_spec(type="image", X=X, steps=10) + + with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir: + rkme_path = os.path.join(tempdir, "rkme.json") + image_rkme.save(rkme_path) + + with open(rkme_path, "r") as f: + data = json.load(f) + assert data["type"] == "RKMEImageSpecification" + + rkme2 = RKMEImageSpecification() + rkme2.load(rkme_path) + assert rkme2.type == "RKMEImageSpecification" + + def test_image_rkme(self): + self._test_image_rkme(np.random.randint(0, 255, size=(2000, 3, 32, 32))) + self._test_image_rkme(np.random.randint(0, 255, size=(100, 1, 128, 128))) + self._test_image_rkme(np.random.randint(0, 255, size=(50, 3, 128, 128)) / 255) + self._test_image_rkme(torch.randint(0, 255, (2000, 3, 32, 32))) + self._test_image_rkme(torch.randint(0, 255, (20, 3, 128, 128))) + self._test_image_rkme(torch.randint(0, 255, (1, 1, 128, 128)) / 255) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_specification/test_rkme.py b/tests/test_specification/test_rkme.py deleted file mode 100644 index 3b33e14..0000000 --- a/tests/test_specification/test_rkme.py +++ /dev/null @@ -1,104 +0,0 @@ -import os -import json -import string -import random -import torch -import unittest -import tempfile -import numpy as np - -from learnware.specification import RKMETableSpecification, RKMEImageSpecification, RKMETextSpecification -from learnware.specification import generate_stat_spec - - -class TestRKME(unittest.TestCase): - def test_rkme(self): - def _test_table_rkme(X): - rkme = generate_stat_spec(type="table", X=X) - - with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir: - rkme_path = os.path.join(tempdir, "rkme.json") - rkme.save(rkme_path) - - with open(rkme_path, "r") as f: - data = json.load(f) - assert data["type"] == "RKMETableSpecification" - - rkme2 = RKMETableSpecification() - rkme2.load(rkme_path) - assert rkme2.type == "RKMETableSpecification" - - _test_table_rkme(np.random.uniform(-10000, 10000, size=(5000, 200))) - _test_table_rkme(np.random.uniform(-10000, 10000, size=(10000, 100))) - _test_table_rkme(np.random.uniform(-10000, 10000, size=(5, 20))) - _test_table_rkme(np.random.uniform(-10000, 10000, size=(1, 50))) - _test_table_rkme(np.random.uniform(-10000, 10000, size=(100, 150))) - - def test_image_rkme(self): - def _test_image_rkme(X): - image_rkme = generate_stat_spec(type="image", X=X, steps=10) - - with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir: - rkme_path = os.path.join(tempdir, "rkme.json") - image_rkme.save(rkme_path) - - with open(rkme_path, "r") as f: - data = json.load(f) - assert data["type"] == "RKMEImageSpecification" - - rkme2 = RKMEImageSpecification() - rkme2.load(rkme_path) - assert rkme2.type == "RKMEImageSpecification" - - _test_image_rkme(np.random.randint(0, 255, size=(2000, 3, 32, 32))) - _test_image_rkme(np.random.randint(0, 255, size=(100, 1, 128, 128))) - _test_image_rkme(np.random.randint(0, 255, size=(50, 3, 128, 128)) / 255) - - _test_image_rkme(torch.randint(0, 255, (2000, 3, 32, 32))) - _test_image_rkme(torch.randint(0, 255, (20, 3, 128, 128))) - _test_image_rkme(torch.randint(0, 255, (1, 1, 128, 128)) / 255) - - def test_text_rkme(self): - def generate_random_text_list(num, text_type="en", min_len=10, max_len=1000): - text_list = [] - for i in range(num): - length = random.randint(min_len, max_len) - if text_type == "en": - characters = string.ascii_letters + string.digits + string.punctuation - result_str = "".join(random.choice(characters) for i in range(length)) - text_list.append(result_str) - elif text_type == "zh": - result_str = "".join(chr(random.randint(0x4E00, 0x9FFF)) for i in range(length)) - text_list.append(result_str) - else: - raise ValueError("Type should be en or zh") - return text_list - - def _test_text_rkme(X): - rkme = generate_stat_spec(type="text", X=X) - - with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir: - rkme_path = os.path.join(tempdir, "rkme.json") - rkme.save(rkme_path) - - with open(rkme_path, "r") as f: - data = json.load(f) - assert data["type"] == "RKMETextSpecification" - - rkme2 = RKMETextSpecification() - rkme2.load(rkme_path) - assert rkme2.type == "RKMETextSpecification" - - return rkme2.get_z().shape[1] - - dim1 = _test_text_rkme(generate_random_text_list(3000, "en")) - dim2 = _test_text_rkme(generate_random_text_list(100, "en")) - dim3 = _test_text_rkme(generate_random_text_list(50, "zh")) - dim4 = _test_text_rkme(generate_random_text_list(5000, "zh")) - dim5 = _test_text_rkme(generate_random_text_list(1, "zh")) - - assert dim1 == dim2 and dim2 == dim3 and dim3 == dim4 and dim4 == dim5 - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_specification/test_table_rkme.py b/tests/test_specification/test_table_rkme.py new file mode 100644 index 0000000..ed26314 --- /dev/null +++ b/tests/test_specification/test_table_rkme.py @@ -0,0 +1,39 @@ +import os +import json +import string +import random +import torch +import unittest +import tempfile +import numpy as np + +from learnware.specification import RKMETableSpecification, RKMEImageSpecification, RKMETextSpecification +from learnware.specification import generate_stat_spec + + +class TestTableRKME(unittest.TestCase): + @staticmethod + def _test_table_rkme(X): + rkme = generate_stat_spec(type="table", X=X) + + with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir: + rkme_path = os.path.join(tempdir, "rkme.json") + rkme.save(rkme_path) + + with open(rkme_path, "r") as f: + data = json.load(f) + assert data["type"] == "RKMETableSpecification" + + rkme2 = RKMETableSpecification() + rkme2.load(rkme_path) + assert rkme2.type == "RKMETableSpecification" + + def test_table_rkme(self): + self._test_table_rkme(np.random.uniform(-10000, 10000, size=(5000, 200))) + self._test_table_rkme(np.random.uniform(-10000, 10000, size=(10000, 100))) + self._test_table_rkme(np.random.uniform(-10000, 10000, size=(5, 20))) + self._test_table_rkme(np.random.uniform(-10000, 10000, size=(1, 50))) + self._test_table_rkme(np.random.uniform(-10000, 10000, size=(100, 150))) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_specification/test_text_rkme.py b/tests/test_specification/test_text_rkme.py new file mode 100644 index 0000000..0409d98 --- /dev/null +++ b/tests/test_specification/test_text_rkme.py @@ -0,0 +1,58 @@ +import os +import json +import string +import random +import unittest +import tempfile + +from learnware.specification import RKMETextSpecification +from learnware.specification import generate_stat_spec + + +class TestTextRKME(unittest.TestCase): + @staticmethod + def generate_random_text_list(num, text_type="en", min_len=10, max_len=1000): + text_list = [] + for i in range(num): + length = random.randint(min_len, max_len) + if text_type == "en": + characters = string.ascii_letters + string.digits + string.punctuation + result_str = "".join(random.choice(characters) for i in range(length)) + text_list.append(result_str) + elif text_type == "zh": + result_str = "".join(chr(random.randint(0x4E00, 0x9FFF)) for i in range(length)) + text_list.append(result_str) + else: + raise ValueError("Type should be en or zh") + return text_list + + @staticmethod + def _test_text_rkme(X): + rkme = generate_stat_spec(type="text", X=X) + + with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir: + rkme_path = os.path.join(tempdir, "rkme.json") + rkme.save(rkme_path) + + with open(rkme_path, "r") as f: + data = json.load(f) + assert data["type"] == "RKMETextSpecification" + + rkme2 = RKMETextSpecification() + rkme2.load(rkme_path) + assert rkme2.type == "RKMETextSpecification" + + return rkme2.get_z().shape[1] + + def test_text_rkme(self): + dim1 = self._test_text_rkme(self.generate_random_text_list(3000, "en")) + dim2 = self._test_text_rkme(self.generate_random_text_list(100, "en")) + dim3 = self._test_text_rkme(self.generate_random_text_list(50, "zh")) + dim4 = self._test_text_rkme(self.generate_random_text_list(5000, "zh")) + dim5 = self._test_text_rkme(self.generate_random_text_list(1, "zh")) + + assert dim1 == dim2 and dim2 == dim3 and dim3 == dim4 and dim4 == dim5 + + +if __name__ == "__main__": + unittest.main()