truefoundry 0.5.1rc3__tar.gz → 0.5.1rc4__tar.gz

{truefoundry-0.5.1rc3 → truefoundry-0.5.1rc4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.5.1rc3
+Version: 0.5.1rc4
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com

{truefoundry-0.5.1rc3 → truefoundry-0.5.1rc4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truefoundry"
-version = "0.5.1rc3"
+version = "0.5.1rc4"
 description = "Truefoundry CLI"
 authors = ["Abhishek Choudhary <abhishek@truefoundry.com>"]
 readme = "README.md"

{truefoundry-0.5.1rc3 → truefoundry-0.5.1rc4}/truefoundry/common/auth_service_client.py

@@ -51,7 +51,7 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
         self._api_server_url = url
 
     def refresh_token(self, token: Token, host: Optional[str] = None) -> Token:
-        host_arg_str = f"--host {host}" if host else "--host HOST"
+        host_arg_str = host if host else "HOST"
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
             raise Exception(
@@ -128,7 +128,7 @@ class AuthServerServiceClient(AuthServiceClient):
         self._auth_server_url = url
 
     def refresh_token(self, token: Token, host: Optional[str] = None) -> Token:
-        host_arg_str = f"--host {host}" if host else "--host HOST"
+        host_arg_str = host if host else "HOST"
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
             raise Exception(

{truefoundry-0.5.1rc3 → truefoundry-0.5.1rc4}/truefoundry/ml/__init__.py

@@ -19,10 +19,7 @@ from truefoundry.ml.log_types.artifacts.artifact import (
     BlobStorageDirectory,
 )
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory, DataDirectoryPath
-from truefoundry.ml.log_types.artifacts.model import (
-    ModelVersion,
-    infer_signature,
-)
+from truefoundry.ml.log_types.artifacts.model import ModelVersion
 from truefoundry.ml.logger import init_logger
 from truefoundry.ml.mlfoundry_api import get_client
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
@@ -42,6 +39,7 @@ from truefoundry.ml.model_framework import (
     TensorFlowFramework,
     TransformersFramework,
     XGBoostFramework,
+    sklearn_infer_schema,
 )
 
 __all__ = [
@@ -80,7 +78,7 @@ __all__ = [
     "XGBoostFramework",
     "XGBoostModelSchema",
     "get_client",
-    "infer_signature",
+    "sklearn_infer_schema",
 ]
 
 init_logger()

{truefoundry-0.5.1rc3 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/models/signature.py

@@ -1,13 +1,16 @@
 from dataclasses import dataclass, is_dataclass
-from typing import Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import numpy as np
-import pandas as pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
 
 from .schema import ParamSchema, Schema, convert_dataclass_to_schema
 from .utils import infer_param_schema, infer_schema
 
-MlflowInferableDataset = Union[pd.DataFrame, np.ndarray, Dict[str, np.ndarray]]
+MlflowInferableDataset = Union["pd.DataFrame", "np.ndarray", Dict[str, "np.ndarray"]]
 
 
 class ModelSignature:

{truefoundry-0.5.1rc3 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/models/utils.py

@@ -4,7 +4,12 @@ from collections import defaultdict
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
-import pandas as pd
+
+try:
+    import pandas as pd
+except ImportError:
+    pd = None
+
 
 from .exceptions import MlflowException
 from .schema import (
@@ -330,7 +335,7 @@ def infer_schema(data: Any) -> Schema:  # noqa: C901
                 ]
             )
     # pandas.Series
-    elif isinstance(data, pd.Series):
+    elif pd and isinstance(data, pd.Series):
         name = getattr(data, "name", None)
         schema = Schema(
             [
@@ -342,7 +347,7 @@ def infer_schema(data: Any) -> Schema:  # noqa: C901
             ]
         )
     # pandas.DataFrame
-    elif isinstance(data, pd.DataFrame):
+    elif pd and isinstance(data, pd.DataFrame):
         schema = Schema(
             [
                 ColSpec(
@@ -473,13 +478,13 @@ def _is_none_or_nan(x):
 
 
 def _infer_required(col) -> bool:
-    if isinstance(col, (list, pd.Series)):
+    if pd and isinstance(col, (list, pd.Series)):
         return not any(_is_none_or_nan(x) for x in col)
     return not _is_none_or_nan(col)
 
 
-def _infer_pandas_column(col: pd.Series) -> DataType:
-    if not isinstance(col, pd.Series):
+def _infer_pandas_column(col: "pd.Series") -> DataType:
+    if pd and not isinstance(col, pd.Series):
         raise TypeError(f"Expected pandas.Series, got '{type(col)}'.")
     if len(col.values.shape) > 1:
         raise MlflowException(f"Expected 1d array, got array with shape {col.shape}")
@@ -496,7 +501,7 @@ def _infer_pandas_column(col: pd.Series) -> DataType:
             # For backwards compatibility, we fall back to string
             # if the provided array is of string type
             # This is for diviner test where df field is ('key2', 'key1', 'key0')
-            if pd.api.types.is_string_dtype(col):
+            if pd and pd.api.types.is_string_dtype(col):
                 return DataType.string
             raise MlflowException(
                 f"Failed to infer schema for pandas.Series {col}. Error: {e}"

{truefoundry-0.5.1rc3 → truefoundry-0.5.1rc4}/truefoundry/ml/log_types/artifacts/model.py

@@ -31,6 +31,7 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     TrueFoundryArtifactSource,
     UpdateModelVersionRequestDto,
 )
+from truefoundry.ml.autogen.models import infer_signature as _infer_signature
 from truefoundry.ml.enums import ModelFramework
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.artifact import BlobStorageDirectory
@@ -611,9 +612,6 @@ def infer_signature(
     ] = None,
     params: Optional[Dict[str, Any]] = None,
 ):
-    # TODO: Importing this globally causes hard dependencies on some libraries like pandas
-    from truefoundry.ml.autogen.models import infer_signature as _infer_signature
-
     return _infer_signature(
         model_input=model_input, model_output=model_output, params=params
     )

{truefoundry-0.5.1rc3 → truefoundry-0.5.1rc4}/truefoundry/ml/model_framework.py

@@ -1,8 +1,20 @@
+import json
 import os
 import warnings
 from collections import OrderedDict
 from pickle import load as pickle_load
-from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union, get_args
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Type,
+    Union,
+    get_args,
+)
 
 from truefoundry.common.utils import (
     get_python_version_major_minor,
@@ -10,6 +22,7 @@ from truefoundry.common.utils import (
 )
 from truefoundry.ml.autogen.client import SerializationFormat
 from truefoundry.ml.autogen.entities import artifacts as autogen_artifacts
+from truefoundry.ml.autogen.models import infer_signature
 from truefoundry.ml.enums import ModelFramework
 from truefoundry.ml.log_types.artifacts.utils import (
     get_single_file_path_if_only_one_in_directory,
@@ -17,6 +30,9 @@ from truefoundry.ml.log_types.artifacts.utils import (
 )
 from truefoundry.pydantic_v1 import BaseModel, Field
 
+if TYPE_CHECKING:
+    from sklearn.base import BaseEstimator
+
 # Map serialization format to corresponding pip packages
 SERIALIZATION_FORMAT_TO_PACKAGES_NAME_MAP = {
     SerializationFormat.JOBLIB: ["joblib"],
@@ -421,3 +437,37 @@ def auto_update_model_framework_details(
                     model_file_path=absolute_model_filepath
                 )
             )
+
+
+def sklearn_infer_schema(
+    model_input: Any,
+    model: "BaseEstimator",
+    infer_method_name: str = "predict",
+) -> autogen_artifacts.SklearnModelSchema:
+    """
+    Infer the schema of a Sklearn model.
+
+    Args:
+        model_input (Any): The input data to be used for schema inference.
+        model (Any): The Sklearn model instance.
+        infer_method_name (str): The name of the method to be used for schema inference.
+            Eg: predict (default), predict_proba
+    Returns:
+        SklearnModelSchema: The inferred schema of the Sklearn model.
+    """
+    if not hasattr(model, infer_method_name):
+        raise ValueError(
+            f"Model does not have the method '{infer_method_name}' to infer the schema."
+        )
+    model_infer_method = getattr(model, infer_method_name)
+    model_output = model_infer_method(model_input)
+
+    model_signature = infer_signature(
+        model_input=model_input, model_output=model_output
+    )
+    model_signature_json = model_signature.to_dict()
+    return autogen_artifacts.SklearnModelSchema(
+        infer_method_name=infer_method_name,
+        inputs=json.loads(model_signature_json["inputs"]),
+        outputs=json.loads(model_signature_json["outputs"]),
+    )