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

atlasai/toolkit/model.py

@@ -12,13 +12,13 @@
 # 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
+from .utils import is_notebook
 
 
 def search(name=None, version=None, source=None, tags=None, aliases=None):
@@ -32,6 +32,12 @@ class ModelPredictWrapper:
     def __init__(self, model):
         self._model = model
 
+    def __repr__(self):
+        return f'Model {self._model.name} - version: {self._model.version}'
+
+    def __str__(self):
+        return f'Model {self._model.name} - version: {self._model.version}'
+
     def predict(
         self,
         data: Union[dict, str, pd.DataFrame] = None,
@@ -67,16 +73,29 @@ class ModelResults:
     def __iter__(self):
         return iter(self.results)
 
+    def __getitem__(self, index):
+        return self.results[index]
+
+    def display_tags(self, tags):
+        if not isinstance(tags, dict):
+            return tags
+        return '\n'.join([f'{k}:{v}' for k, v in tags.items()])
+
     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
+                'tags': self.display_tags(r.tags)
             } for r in self.results
         ]
-        return pd.DataFrame(data)
+        df = pd.DataFrame(data)
+        if is_notebook():
+            from IPython.display import display, HTML
+            html = df.to_html().replace("\\n", "<br>")
+            display(HTML(html))
+        return df
 
     def first(self):
         return self.results[0]
@@ -100,19 +119,19 @@ class ModelSearch:
         self.aliases = aliases if isinstance(aliases, list) else [aliases]
 
     def search(self):
-        search = {}
+        _search = {}
         filters = ['name', 'version', 'source', 'tags', 'aliases']
         for f in filters:
             if not getattr(self, f, None):
                 continue
-            search[f] = getattr(self, f)
+            _search[f] = getattr(self, f)
 
         _get_access_token()
-        results = client.get_models(search=search, tabular=False)
+        results = client.get_models(search=_search, tabular=False)
         return ModelResults([ModelPredictWrapper(r) for r in results])
 
 class ModelGet:
-    def __init__(self, name, version='latest'):
+    def __init__(self, name, version=None):
         self.name = name
         self.version = version
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlasai-dstoolkit-client
-Version: 0.0.10
+Version: 0.0.12
 Summary: UNKNOWN
 Home-page: UNKNOWN
 Author: AtlasAI SWE

{atlasai_dstoolkit_client-0.0.10.dist-info → atlasai_dstoolkit_client-0.0.12.dist-info}/RECORD

@@ -6,11 +6,11 @@ atlasai/toolkit/fabric.py,sha256=6aFR2PGQc9P3Qa07WdBg9eKoUzU8n2y_-gGjYcyMrWY,192
 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/model.py,sha256=RUe0HbDpzvHOV9A4rzG3PgN9boMWDHQ2tR7IKHXzbx8,4126
 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.12.dist-info/LICENSE.txt,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+atlasai_dstoolkit_client-0.0.12.dist-info/METADATA,sha256=2vdA8hWTJEIvSm1OgPRt9wmsTtZO_5oR_znxLDNuDbI,1405
+atlasai_dstoolkit_client-0.0.12.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+atlasai_dstoolkit_client-0.0.12.dist-info/top_level.txt,sha256=HRTbErU8nmHFDaJJ5R_XYbwpt21dqdjDpSva8xyy_0k,8
+atlasai_dstoolkit_client-0.0.12.dist-info/RECORD,,