atlasai-dstoolkit-client 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

atlasai/toolkit/__init__.py

@@ -33,5 +33,5 @@ del osp
 del os
 
 
-from . import fabric, feature
+from . import fabric, feature, dataset, model
 from .login import login, logout

atlasai/toolkit/constants.py

@@ -3,6 +3,9 @@ import os
 VINZ_URL = os.getenv('VINZ_URL') or 'https://vinz.atlasai.co'
 DS_TOOLKIT_URL = os.getenv('DS_TOOLKIT_URL') or 'https://dstoolkit.atlasai.co'
 
+# force mlhub library to go to ds toolkit
+os.environ['MLHUB_URL'] = f'{DS_TOOLKIT_URL}/api'
+
 DEFAULT_PAGE_SIZE = 20
 DISABLE_SSL_VERIFICATION = 'DISABLE_SSL_VERIFICATION'
 

atlasai/toolkit/dataset.py

@@ -0,0 +1,108 @@
+# Copyright 2025 AtlasAI PBC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import uuid
+import io
+import pandas as pd
+
+from . import api, constants, utils
+from .requests import get_session
+
+def register(product_name: str, data=None, file=None, start_date: str = None, end_date: str = None, tags: dict = None, type: str = 'Other'):
+    """
+
+    Parameters
+    ----------
+    data: The data you want to upload
+    file: The file you want to upload. It's either a file or data in RAM
+    product_name: The name of the product you want to be registered
+    start_date: Effective start date of the product in YYYY-MM-DD format
+    end_date: Effective end date of the product in YYYY-MM-DD format
+    tags: Tags you want to add to the product and instance
+    type: Type of the asset
+
+    Returns
+    -------
+
+    """
+    if data is None and file is None:
+        raise Exception('data or file parameter required')
+
+    if file:
+        with open('data.json', 'r', encoding='utf-8') as f:
+            file_contents = f.read()
+            data = io.StringIO(file_contents)
+    else:
+        if isinstance(data, pd.DataFrame):
+            data = data.to_json(orient='records')
+        data = io.StringIO(data)
+
+    additional_params = [
+        ('start_date', start_date),
+        ('end_date', end_date),
+        ('tags', tags),
+        ('type', type)
+    ]
+    _, result = api._post(resource='dataset/upload')
+    result = result['data']
+    session = get_session()
+    response = session.put(result['signed_url'], data=data, headers=result['headers'])
+    response.raise_for_status()
+
+    register_data = {
+        'product_name': product_name,
+        'path': result['storage_path'],
+    }
+    for name, value in additional_params:
+        if value:
+            register_data[name] = value
+
+    _, data = api._post(resource='dataset/register', data=register_data)
+    return {
+        'product_name': data['data']['product_name'],
+        'storage_path': result['storage_path']
+    }
+
+def search(*args, **kwargs):
+    ds = DatasetSearch(*args, **kwargs)
+    utils.show_page(ds.page)
+    return ds
+
+class DatasetSearch:
+
+    def __init__(self, id=None, *args, **kwargs):
+        self.id = id or uuid.uuid4()
+        self._search = None
+
+    def __repr__(self):
+        return f'DatasetSearch({self.id})'
+
+    def __str__(self):
+        return f'DatasetSearch({self.id})'
+
+    @property
+    def page(self):
+        return f'{constants.DS_TOOLKIT_URL}/dataset/search/{self.id}'
+
+    @property
+    def search(self):
+        return self._search
+
+    def info(self):
+        return self._info()
+
+    def _info(self):
+        resource = f'dataset/search/{self.id}/info'
+        _, data = api._get(resource=resource)
+        return data['data']

atlasai/toolkit/model.py

@@ -0,0 +1,121 @@
+# Copyright 2025 AtlasAI PBC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import pandas as pd
+from typing import Union
+
+from atlasai.mlhub import client
+
+from .api import _get_access_token
+
+
+def search(name=None, version=None, source=None, tags=None, aliases=None):
+    return ModelSearch(name, version, source, tags, aliases).search()
+
+def get(name, version=None):
+    return ModelGet(name, version).get()
+
+
+class ModelPredictWrapper:
+    def __init__(self, model):
+        self._model = model
+
+    def predict(
+        self,
+        data: Union[dict, str, pd.DataFrame] = None,
+        deployment_type: str = 'http',
+        timeout: int = 3600,
+        wait_for_completion: bool = True,
+        tabular: bool = True,
+        data_type: str = 'json'
+    ):
+        if deployment_type not in self._model.deployments:
+            raise Exception(f'This model has no {deployment_type} deployment.')
+
+        return client.evaluate(
+            self._model.name,
+            self._model.version,
+            deployment_type=deployment_type,
+            data=data,
+            wait_for_completion=wait_for_completion,
+            timeout=timeout,
+            tabular=tabular,
+            data_type=data_type
+
+        )
+
+    def __getattr__(self, item):
+        return getattr(self._model, item)
+
+
+class ModelResults:
+    def __init__(self, results):
+        self.results = results
+
+    def __iter__(self):
+        return iter(self.results)
+
+    def display(self):
+        data = [
+            {
+                'name': r.name,
+                'version': r.version,
+                'aliases': r.aliases if r.aliases else '',
+                'tags': json.dumps(r.tags) if isinstance(r.tags, dict) else r.tags
+            } for r in self.results
+        ]
+        return pd.DataFrame(data)
+
+    def first(self):
+        return self.results[0]
+
+    def last(self):
+        return self.results[-1]
+
+    def all(self):
+        return self.results
+
+class ModelSearch:
+    def __init__(self, name=None, version=None, source=None, tags=None, aliases=None):
+        if aliases is None:
+            aliases = []
+        if tags is None:
+            tags = {}
+        self.name = name
+        self.version = version
+        self.source = source
+        self.tags = tags
+        self.aliases = aliases if isinstance(aliases, list) else [aliases]
+
+    def search(self):
+        search = {}
+        filters = ['name', 'version', 'source', 'tags', 'aliases']
+        for f in filters:
+            if not getattr(self, f, None):
+                continue
+            search[f] = getattr(self, f)
+
+        _get_access_token()
+        results = client.get_models(search=search, tabular=False)
+        return ModelResults([ModelPredictWrapper(r) for r in results])
+
+class ModelGet:
+    def __init__(self, name, version='latest'):
+        self.name = name
+        self.version = version
+
+    def get(self):
+        result = client.get_model_info(self.name, self.version)
+        return ModelPredictWrapper(result)

{atlasai_dstoolkit_client-0.0.9.dist-info → atlasai_dstoolkit_client-0.0.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlasai-dstoolkit-client
-Version: 0.0.9
+Version: 0.0.10
 Summary: UNKNOWN
 Home-page: UNKNOWN
 Author: AtlasAI SWE
@@ -10,6 +10,7 @@ Platform: UNKNOWN
 Requires-Python: >=3.10
 License-File: LICENSE.txt
 Requires-Dist: arrow (<2.0.0,>=1.3.0)
+Requires-Dist: atlasai-mlhub-client
 Requires-Dist: furl (<3.0.0,>=2.1.2)
 Requires-Dist: pandas (<3.0.0,>=2.2.3)
 Requires-Dist: pyarrow (>=19.0.1)

atlasai_dstoolkit_client-0.0.10.dist-info/RECORD

@@ -0,0 +1,16 @@
+atlasai/toolkit/__init__.py,sha256=rjt3GfV7vA3RFMIzqPwBebJUXdQbeTLmEYmm7SOoOGw,1002
+atlasai/toolkit/api.py,sha256=BvO-gLRmbmkKduwbbADjcLlIkS9blzfM_cbMR4DhQmU,5269
+atlasai/toolkit/constants.py,sha256=Jxozn9tKOvAxyOYOZ6bzFtI9m1YETJF1GURzTl9NNC8,422
+atlasai/toolkit/dataset.py,sha256=fhzhxF9YMzIwEpaJZEPOK2SuLMJhVGI75eahxnH_T2c,3254
+atlasai/toolkit/fabric.py,sha256=6aFR2PGQc9P3Qa07WdBg9eKoUzU8n2y_-gGjYcyMrWY,1921
+atlasai/toolkit/feature.py,sha256=VKwX_yAwwMo1WuLZ_RbKVJEH4L1PAoYEPzWa0lH9ats,2828
+atlasai/toolkit/init.py,sha256=JkdJ6QGdYWrq65jgz2pn5RYXUeUe2Ez88_-eMf5CNi0,1100
+atlasai/toolkit/login.py,sha256=n4ydfo9qCsmbZq6er1xeljBD76vdTJGjbhYHMmOyDbQ,3061
+atlasai/toolkit/model.py,sha256=6gmIVx7joCpVxSQbXlmjkxAjreL_J104laAgKKgGOo0,3525
+atlasai/toolkit/requests.py,sha256=X86nIo07hAjUlilZcZ1lV8RB7KOsTKbTGtcY_SpFEXY,1223
+atlasai/toolkit/utils.py,sha256=lYh3P2XOshRgHCjFeXJ0FOJWQW64sddgx8c2kL6Wqwc,1566
+atlasai_dstoolkit_client-0.0.10.dist-info/LICENSE.txt,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+atlasai_dstoolkit_client-0.0.10.dist-info/METADATA,sha256=h9oUexHohAl4WiRMqQX2UiXytdzdYmBeGPJ5BhM4thg,1405
+atlasai_dstoolkit_client-0.0.10.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+atlasai_dstoolkit_client-0.0.10.dist-info/top_level.txt,sha256=HRTbErU8nmHFDaJJ5R_XYbwpt21dqdjDpSva8xyy_0k,8
+atlasai_dstoolkit_client-0.0.10.dist-info/RECORD,,

atlasai_dstoolkit_client-0.0.9.dist-info/RECORD

@@ -1,14 +0,0 @@
-atlasai/toolkit/__init__.py,sha256=szu5LHZ41ccWztCI2LvQ3QnZRybosgTbU3yd6tS7cUw,986
-atlasai/toolkit/api.py,sha256=BvO-gLRmbmkKduwbbADjcLlIkS9blzfM_cbMR4DhQmU,5269
-atlasai/toolkit/constants.py,sha256=sE0PeFa9_htCPVFADHbkyPIz3QOai0tAIA5IiiZI8wA,329
-atlasai/toolkit/fabric.py,sha256=6aFR2PGQc9P3Qa07WdBg9eKoUzU8n2y_-gGjYcyMrWY,1921
-atlasai/toolkit/feature.py,sha256=VKwX_yAwwMo1WuLZ_RbKVJEH4L1PAoYEPzWa0lH9ats,2828
-atlasai/toolkit/init.py,sha256=JkdJ6QGdYWrq65jgz2pn5RYXUeUe2Ez88_-eMf5CNi0,1100
-atlasai/toolkit/login.py,sha256=n4ydfo9qCsmbZq6er1xeljBD76vdTJGjbhYHMmOyDbQ,3061
-atlasai/toolkit/requests.py,sha256=X86nIo07hAjUlilZcZ1lV8RB7KOsTKbTGtcY_SpFEXY,1223
-atlasai/toolkit/utils.py,sha256=lYh3P2XOshRgHCjFeXJ0FOJWQW64sddgx8c2kL6Wqwc,1566
-atlasai_dstoolkit_client-0.0.9.dist-info/LICENSE.txt,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-atlasai_dstoolkit_client-0.0.9.dist-info/METADATA,sha256=ygcwrboXNHtv0Y4o7PrEcs2RZDcjkFfQwp357F8mIqg,1368
-atlasai_dstoolkit_client-0.0.9.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-atlasai_dstoolkit_client-0.0.9.dist-info/top_level.txt,sha256=HRTbErU8nmHFDaJJ5R_XYbwpt21dqdjDpSva8xyy_0k,8
-atlasai_dstoolkit_client-0.0.9.dist-info/RECORD,,