Modality

aixplain/factories/model_factory.py

@@ -304,10 +304,10 @@ def list_functions(cls, verbose: Optional[bool] = False, api_key: Optional[Text]
     def create_asset_repo(
         cls,
         name: Text,
-        hosting_machine: Text,
-        version: Text,
         description: Text,
         function: Text,
+        input_modality: Text,
+        output_modality: Text,
         source_language: Text,
         api_key: Optional[Text] = None,
     ) -> Dict:
@@ -316,12 +316,10 @@ def create_asset_repo(
 
         Args:
             name (Text): Model name
-            hosting_machine (Text): Hosting machine ID obtained via list_host_machines
-            always_on (bool): Whether the model should always be on
-            version (Text): Model version
             description (Text): Model description
-            function (Text): Model function name obtained via LIST_HOST_MACHINES
-            is_async (bool): Whether the model is asynchronous or not (False in first release)
+            function (Text): Model function name obtained via LIST_FUNCTIONS
+            input_modality (Text): Modality of the inputs to this model.
+            output_modality (Text): Modality of the outputs to this model.
             source_language (Text): 2-character 639-1 code or 3-character 639-3 language code.
             api_key (Text, optional): Team API key. Defaults to None.
 
@@ -336,23 +334,28 @@ def create_asset_repo(
                 function_id = function_dict["id"]
         if function_id is None:
             raise Exception("Invalid function name")
-        create_url = urljoin(config.BACKEND_URL, f"sdk/models/register")
+        create_url = urljoin(config.BACKEND_URL, f"sdk/models/onboard")
         logging.debug(f"URL: {create_url}")
         if api_key:
             headers = {"x-api-key": f"{api_key}", "Content-Type": "application/json"}
         else:
             headers = {"x-api-key": f"{config.TEAM_API_KEY}", "Content-Type": "application/json"}
-        always_on = False
-        is_async = False  # Hard-coded to False for first release
         payload = {
-            "name": name,
-            "hostingMachine": hosting_machine,
-            "alwaysOn": always_on,
-            "version": version,
-            "description": description,
-            "function": function_id,
-            "isAsync": is_async,
-            "sourceLanguage": source_language,
+            "model": {
+                "name": name,
+                "description": description,
+                "connectionType": [
+                    "asynchronous"
+                ],
+                "function": function_id,
+                "modalities": [
+                    f"{input_modality}-{output_modality}"
+                ],
+                "documentationUrl": "aiXplain",
+                "sourceLanguage": source_language
+            },
+            "source": "aixplain-ecr",
+            "onboardingParams": {}
         }
         payload = json.dumps(payload)
         logging.debug(f"Body: {str(payload)}")

tests/functional/model/image_upload_e2e_test.py

@@ -0,0 +1,63 @@
+__author__ = "michaellam"
+
+from pathlib import Path
+import json
+from aixplain.factories.model_factory import ModelFactory
+from tests.test_utils import delete_asset, delete_service_account
+from aixplain.utils import config
+import docker
+import os
+
+def test_create_and_upload_model():
+    # List the functions
+    response = ModelFactory.list_functions()
+    items = response["items"]
+    for item in items:
+        assert "output" not in item.keys()
+        assert "params" not in item.keys()
+        assert "id" not in item.keys()
+        assert "name" in item.keys()
+
+    # Register the model, and create an image repository for it.
+    with open(Path("tests/test_requests/create_asset_request.json")) as f:
+        register_payload_dict = json.load(f)
+    for key in register_payload_dict.keys():
+        register_payload = register_payload_dict[key]
+        name = register_payload["name"]
+        description = register_payload["description"]
+        function = register_payload["function"]
+        input_modality = register_payload["input_modality"]
+        output_modality = register_payload["input_modality"]
+        source_language = register_payload["sourceLanguage"]
+        register_response = ModelFactory.create_asset_repo(name, description, function, input_modality, output_modality, source_language)
+        assert "id" in register_response.keys()
+        assert "repositoryName" in register_response.keys()
+        model_id = register_response["id"]
+        repo_name = register_response["repositoryName"]
+
+        # Log into the image repository.
+        login_response = ModelFactory.asset_repo_login()
+
+        assert login_response["username"] == "AWS"
+        assert login_response["registry"] == "535945872701.dkr.ecr.us-east-1.amazonaws.com"
+        assert "password" in login_response.keys()
+
+        username = login_response["username"]
+        password = login_response["password"]
+        registry = login_response["registry"]
+
+        # Push an image to ECR
+        # os.system("aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 535945872701.dkr.ecr.us-east-1.amazonaws.com")
+        low_level_client = docker.APIClient(base_url='unix://var/run/docker.sock')
+        # low_level_client.pull("535945872701.dkr.ecr.us-east-1.amazonaws.com/bash")
+        # low_level_client.tag("535945872701.dkr.ecr.us-east-1.amazonaws.com/bash", f"{registry}/{repo_name}")
+        low_level_client.pull("bash")
+        low_level_client.tag("bash", f"{registry}/{repo_name}")
+        low_level_client.push(f"{registry}/{repo_name}", auth_config={"username":username, "password":password})
+
+        # Send an email to finalize onboarding process
+        ModelFactory.onboard_model(model_id, "latest", "fake_hash")
+
+        # Clean up
+        delete_service_account(config.TEAM_API_KEY)
+        delete_asset(model_id, config.TEAM_API_KEY)

tests/image_upload_functional_test.py → tests/functional/model/image_upload_functional_test.py

@@ -1,9 +1,10 @@
 __author__ = "michaellam"
+
 from pathlib import Path
 import json
-from aixplain.utils import config
-from tests.test_utils import delete_asset, delete_service_account
 from aixplain.factories.model_factory import ModelFactory
+from tests.test_utils import delete_asset, delete_service_account
+from aixplain.utils import config
 import pytest
 
 def test_login():
@@ -17,21 +18,23 @@ def test_login():
 
 def test_create_asset_repo():
     with open(Path("tests/test_requests/create_asset_request.json")) as f:
-        mock_register_payload = json.load(f)
-    name = mock_register_payload["name"]
-    host_machine = mock_register_payload["hostingMachine"]
-    version = mock_register_payload["version"]
-    description = mock_register_payload["description"]
-    function = mock_register_payload["function"]
-    source_language = mock_register_payload["sourceLanguage"]
-    response = ModelFactory.create_asset_repo(name, host_machine, version, description, function, source_language)
-    print(response)
-    response_dict = dict(response)
-    assert "id" in response_dict.keys()
-    assert "repositoryName" in response_dict.keys()
+        mock_register_payload_dict = json.load(f)
+    for key in mock_register_payload_dict.keys():
+        mock_register_payload = mock_register_payload_dict[key]
+        name = mock_register_payload["name"]
+        description = mock_register_payload["description"]
+        function = mock_register_payload["function"]
+        input_modality = mock_register_payload["input_modality"]
+        output_modality = mock_register_payload["input_modality"]
+        source_language = mock_register_payload["sourceLanguage"]
+        response = ModelFactory.create_asset_repo(name, description, function, input_modality, output_modality, source_language)
+        print(response)
+        response_dict = dict(response)
+        assert "id" in response_dict.keys()
+        assert "repositoryName" in response_dict.keys()
 
-    # Test cleanup
-    delete_asset(response["id"], config.TEAM_API_KEY)
+        # Test cleanup
+        delete_asset(response["id"], config.TEAM_API_KEY)
 
 def test_list_host_machines():
     response = ModelFactory.list_host_machines()

tests/test_requests/create_asset_request.json

@@ -1,8 +1,74 @@
 {
-    "name": "mock_name",
-    "hostingMachine": "aix-2c-8g-od",
-    "version": "mock_version",
-    "description": "mock_description",
-    "function": "Speech Recognition",
-    "sourceLanguage": "en"
+    "text_generation":
+    {
+        "name": "mock_name",
+        "description": "mock_description",
+        "function": "Text Generation",
+        "input_modality": "text",
+        "output_modality": "text",
+        "sourceLanguage": "en"
+    },
+    "text_summarization":
+    {
+        "name": "mock_name",
+        "description": "mock_description",
+        "function": "Text summarization",
+        "input_modality": "text",
+        "output_modality": "text",
+        "sourceLanguage": "en"
+    },
+    "translation":
+    {
+        "name": "mock_name",
+        "description": "mock_description",
+        "function": "Translation",
+        "input_modality": "text",
+        "output_modality": "text",
+        "sourceLanguage": "en"
+    },
+    "text_reconstruction":
+    {
+        "name": "mock_name",
+        "description": "mock_description",
+        "function": "Text Reconstruction",
+        "input_modality": "text",
+        "output_modality": "text",
+        "sourceLanguage": "en"
+    },
+    "subtitling_translation":
+    {
+        "name": "mock_name",
+        "description": "mock_description",
+        "function": "Subtitling Translation",
+        "input_modality": "text",
+        "output_modality": "text",
+        "sourceLanguage": "en"
+    },
+    "fill_text_mask":
+    {
+        "name": "mock_name",
+        "description": "mock_description",
+        "function": "Fill Text Mask",
+        "input_modality": "text",
+        "output_modality": "text",
+        "sourceLanguage": "en"
+    },
+    "diacritization":
+    {
+        "name": "mock_name",
+        "description": "mock_description",
+        "function": "Diacritization",
+        "input_modality": "text",
+        "output_modality": "text",
+        "sourceLanguage": "en"
+    },
+    "search":
+    {
+        "name": "mock_name",
+        "description": "mock",
+        "function": "Search",
+        "input_modality": "text",
+        "output_modality": "text",
+        "sourceLanguage": "en"
+    }
 }

tests/image_upload_test.py → tests/unit/image_upload_test.py

@@ -23,7 +23,7 @@ def test_login():
     assert creds == mock_json
 
 def test_create_asset_repo():
-    url_register =  urljoin(config.BACKEND_URL, f"sdk/models/register")
+    url_register =  urljoin(config.BACKEND_URL, f"sdk/models/onboard")
     url_function = urljoin(config.BACKEND_URL, f"sdk/functions")
     with requests_mock.Mocker() as mock:
         with open(Path("tests/mock_responses/create_asset_repo_response.json")) as f:
@@ -32,8 +32,13 @@ def test_create_asset_repo():
         with open(Path("tests/mock_responses/list_functions_response.json")) as f:
             mock_json_functions = json.load(f)
         mock.get(url_function, headers=AUTH_FIXED_HEADER, json=mock_json_functions)
-        model_id = ModelFactory.create_asset_repo("mock_name", "mock_machines", "mock_version", 
-                          "mock_description", "Speech Recognition", "en", config.TEAM_API_KEY)
+        model_id = ModelFactory.create_asset_repo("mock_name", 
+                                                  "mock_description", 
+                                                  "Text Generation", 
+                                                  "text", 
+                                                  "text", 
+                                                  "en", 
+                                                  config.TEAM_API_KEY)
     assert model_id == mock_json_register
 
 def test_list_host_machines():