Rename base client to state client

NillionNetwork · Jun 7, 2024 · 2348884 · 2348884
1 parent 6739dcc
commit 2348884
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/nada_ai/__init__.py b/nada_ai/__init__.py
@@ -1 +1 @@
-from nada_ai.client import BaseClient, TorchClient, SklearnClient
+from nada_ai.client import StateClient, TorchClient, SklearnClient
diff --git a/nada_ai/client.py b/nada_ai/client.py
@@ -106,8 +106,8 @@ def __ensure_numpy(self, array_like: Any) -> np.ndarray:
         )
 
 
-class BaseClient(ModelClient):
-    """Base ModelClient for generic model states"""
+class StateClient(ModelClient):
+    """ModelClient for generic model states"""
 
     def __init__(self, state_dict: Dict[str, Any]) -> None:
         """

diff --git a/tests/python-tests/test_model_client.py b/tests/python-tests/test_model_client.py
@@ -8,7 +8,7 @@
 from sklearn.linear_model import LinearRegression, LogisticRegression
 from torch import nn
 from nada_ai.client import ModelClient
-from nada_ai import BaseClient, TorchClient, SklearnClient
+from nada_ai import StateClient, TorchClient, SklearnClient
 import py_nillion_client as nillion
 
 
@@ -73,10 +73,10 @@ def test_sklearn_4(self):
         with pytest.raises(NotImplementedError):
             model_client.export_state_as_secrets("test_model", nillion.SecretInteger)
 
-    def test_base_client_1(self):
-        base_client = BaseClient({"some_value": 1})
+    def test_state_client_1(self):
+        state_client = StateClient({"some_value": 1})
 
-        secrets = base_client.export_state_as_secrets("test_model", na.Rational)
+        secrets = state_client.export_state_as_secrets("test_model", na.Rational)
 
         assert list(secrets.keys()) == ["test_model_some_value_0"]