truefoundry 0.5.1rc1__py3-none-any.whl → 0.5.1rc3__py3-none-any.whl

truefoundry/ml/__init__.py

@@ -1,5 +1,9 @@
-from truefoundry.ml.autogen.client.models import LibraryName, ModelVersionEnvironment
-from truefoundry.ml.autogen.models import infer_signature
+from truefoundry.ml.autogen.client.models import (
+    LibraryName,
+    ModelVersionEnvironment,
+    SklearnModelSchema,
+    XGBoostModelSchema,
+)
 from truefoundry.ml.enums import (
     DataSlice,
     FileFormat,
@@ -17,6 +21,7 @@ from truefoundry.ml.log_types.artifacts.artifact import (
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory, DataDirectoryPath
 from truefoundry.ml.log_types.artifacts.model import (
     ModelVersion,
+    infer_signature,
 )
 from truefoundry.ml.logger import init_logger
 from truefoundry.ml.mlfoundry_api import get_client
@@ -66,12 +71,14 @@ __all__ = [
     "Plot",
     "PyTorchFramework",
     "SklearnFramework",
+    "SklearnModelSchema",
     "SpaCyFramework",
     "StatsModelsFramework",
     "TensorFlowFramework",
     "TransformersFramework",
     "ViewType",
     "XGBoostFramework",
+    "XGBoostModelSchema",
     "get_client",
     "infer_signature",
 ]

truefoundry/ml/autogen/client/__init__.py

@@ -372,6 +372,7 @@ from truefoundry.ml.autogen.client.models.set_experiment_tag_request_dto import
 from truefoundry.ml.autogen.client.models.set_tag_request_dto import SetTagRequestDto
 from truefoundry.ml.autogen.client.models.signed_url_dto import SignedURLDto
 from truefoundry.ml.autogen.client.models.sklearn_framework import SklearnFramework
+from truefoundry.ml.autogen.client.models.sklearn_model_schema import SklearnModelSchema
 from truefoundry.ml.autogen.client.models.source import Source
 from truefoundry.ml.autogen.client.models.source1 import Source1
 from truefoundry.ml.autogen.client.models.spa_cy_framework import SpaCyFramework
@@ -427,3 +428,6 @@ from truefoundry.ml.autogen.client.models.validation_error_loc_inner import (
     ValidationErrorLocInner,
 )
 from truefoundry.ml.autogen.client.models.xg_boost_framework import XGBoostFramework
+from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
+    XGBoostModelSchema,
+)

truefoundry/ml/autogen/client/models/__init__.py

@@ -343,6 +343,7 @@ from truefoundry.ml.autogen.client.models.set_experiment_tag_request_dto import
 from truefoundry.ml.autogen.client.models.set_tag_request_dto import SetTagRequestDto
 from truefoundry.ml.autogen.client.models.signed_url_dto import SignedURLDto
 from truefoundry.ml.autogen.client.models.sklearn_framework import SklearnFramework
+from truefoundry.ml.autogen.client.models.sklearn_model_schema import SklearnModelSchema
 from truefoundry.ml.autogen.client.models.source import Source
 from truefoundry.ml.autogen.client.models.source1 import Source1
 from truefoundry.ml.autogen.client.models.spa_cy_framework import SpaCyFramework
@@ -398,3 +399,6 @@ from truefoundry.ml.autogen.client.models.validation_error_loc_inner import (
     ValidationErrorLocInner,
 )
 from truefoundry.ml.autogen.client.models.xg_boost_framework import XGBoostFramework
+from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
+    XGBoostModelSchema,
+)

truefoundry/ml/autogen/client/models/infer_method_name.py

@@ -19,7 +19,7 @@ from aenum import Enum
 
 class InferMethodName(str, Enum):
     """
-    +label=Inference method name +usage=Name of the method used for inference
+    +label=Inference Method Name +usage=Name of the method used for inference
     """
 
     """

truefoundry/ml/autogen/client/models/model_version_manifest.py

@@ -53,10 +53,6 @@ class ModelVersionManifest(BaseModel):
     step: Optional[conint(strict=True, ge=0)] = Field(
         default=0, description="+label=Step"
     )
-    model_schema: Optional[Dict[str, Any]] = Field(
-        default=None,
-        description="+label=Model Schema +usage=Schema of the model +uiType=JsonInput",
-    )
     __properties = [
         "description",
         "metadata",
@@ -65,7 +61,6 @@ class ModelVersionManifest(BaseModel):
         "framework",
         "environment",
         "step",
-        "model_schema",
     ]
 
     @validator("type")
@@ -137,7 +132,6 @@ class ModelVersionManifest(BaseModel):
                 if obj.get("environment") is not None
                 else None,
                 "step": obj.get("step") if obj.get("step") is not None else 0,
-                "model_schema": obj.get("model_schema"),
             }
         )
         return _obj

truefoundry/ml/autogen/client/models/sklearn_framework.py

@@ -18,10 +18,10 @@ import pprint
 import re  # noqa: F401
 from typing import Optional
 
-from truefoundry.ml.autogen.client.models.infer_method_name import InferMethodName
 from truefoundry.ml.autogen.client.models.serialization_format import (
     SerializationFormat,
 )
+from truefoundry.ml.autogen.client.models.sklearn_model_schema import SklearnModelSchema
 from truefoundry.pydantic_v1 import BaseModel, Field, StrictStr, validator
 
 
@@ -34,24 +34,16 @@ class SklearnFramework(BaseModel):
         default=...,
         description="+label=Type +usage=Type of the framework +value=sklearn",
     )
-    serialization_format: Optional[SerializationFormat] = Field(
-        default=None,
-        description="+label=Serialization format +usage=Serialization format used for the model",
-    )
     model_filepath: Optional[StrictStr] = Field(
         default=None,
         description="+label=Model file path +usage=Relative path to the model file",
     )
-    infer_method_name: Optional[InferMethodName] = Field(
+    serialization_format: Optional[SerializationFormat] = Field(
         default=None,
-        description="+label=Inference method name +usage=Name of the method used for inference",
+        description="+label=Serialization format +usage=Serialization format used for the model",
     )
-    __properties = [
-        "type",
-        "serialization_format",
-        "model_filepath",
-        "infer_method_name",
-    ]
+    model_schema: Optional[SklearnModelSchema] = None
+    __properties = ["type", "model_filepath", "serialization_format", "model_schema"]
 
     @validator("type")
     def type_validate_enum(cls, value):
@@ -82,6 +74,9 @@ class SklearnFramework(BaseModel):
     def to_dict(self):
         """Returns the dictionary representation of the model using alias"""
         _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        # override the default output from truefoundry.pydantic_v1 by calling `to_dict()` of model_schema
+        if self.model_schema:
+            _dict["model_schema"] = self.model_schema.to_dict()
         return _dict
 
     @classmethod
@@ -96,9 +91,11 @@ class SklearnFramework(BaseModel):
         _obj = SklearnFramework.parse_obj(
             {
                 "type": obj.get("type"),
-                "serialization_format": obj.get("serialization_format"),
                 "model_filepath": obj.get("model_filepath"),
-                "infer_method_name": obj.get("infer_method_name"),
+                "serialization_format": obj.get("serialization_format"),
+                "model_schema": SklearnModelSchema.from_dict(obj.get("model_schema"))
+                if obj.get("model_schema") is not None
+                else None,
             }
         )
         return _obj

truefoundry/ml/autogen/client/models/sklearn_model_schema.py

@@ -0,0 +1,82 @@
+# coding: utf-8
+
+"""
+FastAPI
+
+No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
+
+The version of the OpenAPI document: 0.1.0
+Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+Do not edit the class manually.
+"""  # noqa: E501
+
+from __future__ import annotations
+
+import json
+import pprint
+import re  # noqa: F401
+from typing import Any, Dict
+
+from truefoundry.ml.autogen.client.models.infer_method_name import InferMethodName
+from truefoundry.pydantic_v1 import BaseModel, Field, conlist
+
+
+class SklearnModelSchema(BaseModel):
+    """
+    SklearnModelSchema
+    """
+
+    infer_method_name: InferMethodName = Field(
+        default=...,
+        description="+label=Inference Method Name +usage=Name of the method used for inference",
+    )
+    inputs: conlist(Dict[str, Any]) = Field(
+        default=..., description="+label= Input Schema +usage=Schema of the input"
+    )
+    outputs: conlist(Dict[str, Any]) = Field(
+        default=..., description="+label= Output Schema +usage=Schema of the output"
+    )
+    __properties = ["infer_method_name", "inputs", "outputs"]
+
+    class Config:
+        """Pydantic configuration"""
+
+        allow_population_by_field_name = True
+        validate_assignment = True
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.dict(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> SklearnModelSchema:
+        """Create an instance of SklearnModelSchema from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self):
+        """Returns the dictionary representation of the model using alias"""
+        _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> SklearnModelSchema:
+        """Create an instance of SklearnModelSchema from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return SklearnModelSchema.parse_obj(obj)
+
+        _obj = SklearnModelSchema.parse_obj(
+            {
+                "infer_method_name": obj.get("infer_method_name"),
+                "inputs": obj.get("inputs"),
+                "outputs": obj.get("outputs"),
+            }
+        )
+        return _obj

truefoundry/ml/autogen/client/models/xg_boost_framework.py

@@ -21,6 +21,9 @@ from typing import Optional
 from truefoundry.ml.autogen.client.models.serialization_format import (
     SerializationFormat,
 )
+from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
+    XGBoostModelSchema,
+)
 from truefoundry.pydantic_v1 import BaseModel, Field, StrictStr, validator
 
 
@@ -41,7 +44,8 @@ class XGBoostFramework(BaseModel):
         default=None,
         description="+label=Model file path +usage=Relative path to the model file",
     )
-    __properties = ["type", "serialization_format", "model_filepath"]
+    model_schema: Optional[XGBoostModelSchema] = None
+    __properties = ["type", "serialization_format", "model_filepath", "model_schema"]
 
     @validator("type")
     def type_validate_enum(cls, value):
@@ -72,6 +76,9 @@ class XGBoostFramework(BaseModel):
     def to_dict(self):
         """Returns the dictionary representation of the model using alias"""
         _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        # override the default output from truefoundry.pydantic_v1 by calling `to_dict()` of model_schema
+        if self.model_schema:
+            _dict["model_schema"] = self.model_schema.to_dict()
         return _dict
 
     @classmethod
@@ -88,6 +95,9 @@ class XGBoostFramework(BaseModel):
                 "type": obj.get("type"),
                 "serialization_format": obj.get("serialization_format"),
                 "model_filepath": obj.get("model_filepath"),
+                "model_schema": XGBoostModelSchema.from_dict(obj.get("model_schema"))
+                if obj.get("model_schema") is not None
+                else None,
             }
         )
         return _obj

truefoundry/ml/autogen/client/models/xg_boost_model_schema.py

@@ -0,0 +1,82 @@
+# coding: utf-8
+
+"""
+FastAPI
+
+No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
+
+The version of the OpenAPI document: 0.1.0
+Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+Do not edit the class manually.
+"""  # noqa: E501
+
+from __future__ import annotations
+
+import json
+import pprint
+import re  # noqa: F401
+from typing import Any, Dict
+
+from truefoundry.ml.autogen.client.models.infer_method_name import InferMethodName
+from truefoundry.pydantic_v1 import BaseModel, Field, conlist
+
+
+class XGBoostModelSchema(BaseModel):
+    """
+    XGBoostModelSchema
+    """
+
+    infer_method_name: InferMethodName = Field(
+        default=...,
+        description="+label=Inference Method Name +usage=Name of the method used for inference",
+    )
+    inputs: conlist(Dict[str, Any]) = Field(
+        default=..., description="+label= Input Schema +usage=Schema of the input"
+    )
+    outputs: conlist(Dict[str, Any]) = Field(
+        default=..., description="+label= Output Schema +usage=Schema of the output"
+    )
+    __properties = ["infer_method_name", "inputs", "outputs"]
+
+    class Config:
+        """Pydantic configuration"""
+
+        allow_population_by_field_name = True
+        validate_assignment = True
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.dict(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> XGBoostModelSchema:
+        """Create an instance of XGBoostModelSchema from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self):
+        """Returns the dictionary representation of the model using alias"""
+        _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> XGBoostModelSchema:
+        """Create an instance of XGBoostModelSchema from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return XGBoostModelSchema.parse_obj(obj)
+
+        _obj = XGBoostModelSchema.parse_obj(
+            {
+                "infer_method_name": obj.get("infer_method_name"),
+                "inputs": obj.get("inputs"),
+                "outputs": obj.get("outputs"),
+            }
+        )
+        return _obj

truefoundry/ml/autogen/client_README.md

@@ -304,6 +304,7 @@ Class | Method | HTTP request | Description
  - [SetTagRequestDto](truefoundry/ml/autogen/client/docs/SetTagRequestDto.md)
  - [SignedURLDto](truefoundry/ml/autogen/client/docs/SignedURLDto.md)
  - [SklearnFramework](truefoundry/ml/autogen/client/docs/SklearnFramework.md)
+ - [SklearnModelSchema](truefoundry/ml/autogen/client/docs/SklearnModelSchema.md)
  - [Source](truefoundry/ml/autogen/client/docs/Source.md)
  - [Source1](truefoundry/ml/autogen/client/docs/Source1.md)
  - [SpaCyFramework](truefoundry/ml/autogen/client/docs/SpaCyFramework.md)
@@ -331,6 +332,7 @@ Class | Method | HTTP request | Description
  - [ValidationError](truefoundry/ml/autogen/client/docs/ValidationError.md)
  - [ValidationErrorLocInner](truefoundry/ml/autogen/client/docs/ValidationErrorLocInner.md)
  - [XGBoostFramework](truefoundry/ml/autogen/client/docs/XGBoostFramework.md)
+ - [XGBoostModelSchema](truefoundry/ml/autogen/client/docs/XGBoostModelSchema.md)
 
 
 <a id="documentation-for-authorization"></a>

truefoundry/ml/autogen/entities/artifacts.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  artifacts.json
-#   timestamp: 2024-12-04T12:00:28+00:00
+#   timestamp: 2024-12-05T14:45:34+00:00
 
 from __future__ import annotations
 
@@ -60,6 +60,19 @@ class BaseArtifactVersion(BaseModel):
     )
 
 
+class BaseModelSchema(BaseModel):
+    infer_method_name: str = Field(
+        ...,
+        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
+    )
+    inputs: List[Dict[str, Any]] = Field(
+        ..., description="+label= Input Schema\n+usage=Schema of the input"
+    )
+    outputs: List[Dict[str, Any]] = Field(
+        ..., description="+label= Output Schema\n+usage=Schema of the output"
+    )
+
+
 class MimeType(str, Enum):
     """
     +label=MIME Type
@@ -289,7 +302,7 @@ class SerializationFormat(str, Enum):
 
 class InferMethodName(str, Enum):
     """
-    +label=Inference method name
+    +label=Inference Method Name
     +usage=Name of the method used for inference
     """
 
@@ -297,26 +310,10 @@ class InferMethodName(str, Enum):
     predict_proba = "predict_proba"
 
 
-class SklearnFramework(BaseModel):
-    """
-    +docs=Scikit-learn framework for the model version
-    +label=Sklearn
-    """
-
-    type: Literal["sklearn"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=sklearn"
-    )
-    serialization_format: Optional[SerializationFormat] = Field(
-        None,
-        description="+label=Serialization format\n+usage=Serialization format used for the model",
-    )
-    model_filepath: Optional[str] = Field(
-        None,
-        description="+label=Model file path\n+usage=Relative path to the model file",
-    )
-    infer_method_name: Optional[InferMethodName] = Field(
-        None,
-        description="+label=Inference method name\n+usage=Name of the method used for inference",
+class SklearnModelSchema(BaseModelSchema):
+    infer_method_name: InferMethodName = Field(
+        ...,
+        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
     )
 
 
@@ -451,22 +448,10 @@ class UserMessage(BaseModel):
     )
 
 
-class XGBoostFramework(BaseModel):
-    """
-    +docs=XGBoost framework for the model version
-    +label=XGBoost
-    """
-
-    type: Literal["xgboost"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=xgboost"
-    )
-    serialization_format: Optional[SerializationFormat] = Field(
-        None,
-        description="+label=Serialization format\n+usage=Serialization format used for the model",
-    )
-    model_filepath: Optional[str] = Field(
-        None,
-        description="+label=Model file path\n+usage=Relative path to the model file",
+class XGBoostModelSchema(BaseModelSchema):
+    infer_method_name: InferMethodName = Field(
+        ...,
+        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
     )
 
 
@@ -563,6 +548,53 @@ class ChatPrompt(BasePrompt):
     )
 
 
+class SklearnFramework(BaseModel):
+    """
+    +docs=Scikit-learn framework for the model version
+    +label=Sklearn
+    """
+
+    type: Literal["sklearn"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=sklearn"
+    )
+    model_filepath: Optional[str] = Field(
+        None,
+        description="+label=Model file path\n+usage=Relative path to the model file",
+    )
+    serialization_format: Optional[SerializationFormat] = Field(
+        None,
+        description="+label=Serialization format\n+usage=Serialization format used for the model",
+    )
+    model_schema: Optional[SklearnModelSchema] = None
+
+
+class XGBoostFramework(BaseModel):
+    """
+    +docs=XGBoost framework for the model version
+    +label=XGBoost
+    """
+
+    type: Literal["xgboost"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=xgboost"
+    )
+    serialization_format: Optional[SerializationFormat] = Field(
+        None,
+        description="+label=Serialization format\n+usage=Serialization format used for the model",
+    )
+    model_filepath: Optional[str] = Field(
+        None,
+        description="+label=Model file path\n+usage=Relative path to the model file",
+    )
+    model_schema: Optional[XGBoostModelSchema] = None
+
+
+class AgentApp(BaseModel):
+    type: Literal["agent-app"] = Field(..., description="+value=agent-app")
+    tools: List[AgentOpenAPIToolWithFQN]
+    agents: List[AgentWithFQN]
+    root_agent: constr(min_length=1)
+
+
 class ModelVersion(BaseArtifactVersion):
     type: Literal["model-version"] = Field(
         ..., description='+label=Type\n+usage=Model Version\n+value="model-version"'
@@ -593,17 +625,6 @@ class ModelVersion(BaseArtifactVersion):
     )
     environment: Optional[ModelVersionEnvironment] = None
     step: conint(ge=0) = Field(0, description="+label=Step")
-    model_schema: Optional[Dict[str, Any]] = Field(
-        None,
-        description="+label=Model Schema\n+usage=Schema of the model\n+uiType=JsonInput",
-    )
-
-
-class AgentApp(BaseModel):
-    type: Literal["agent-app"] = Field(..., description="+value=agent-app")
-    tools: List[AgentOpenAPIToolWithFQN]
-    agents: List[AgentWithFQN]
-    root_agent: constr(min_length=1)
 
 
 class VersionedArtifactType(BaseModel):

truefoundry/ml/log_types/artifacts/model.py

@@ -54,6 +54,9 @@ from truefoundry.ml.session import _get_api_client
 from truefoundry.pydantic_v1 import BaseModel, Extra
 
 if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
 
 
@@ -102,7 +105,6 @@ class ModelVersion:
         self._deleted = False
         self._description: str = ""
         self._metadata: Dict[str, Any] = {}
-        self._model_schema: Optional[Dict[str, Any]] = None
         self._environment: Optional[ModelVersionEnvironment] = None
         self._framework: Optional[ModelFrameworkType] = None
         self._set_mutable_attrs()
@@ -146,9 +148,6 @@ class ModelVersion:
         if self._model_version.manifest:
             self._description = self._model_version.manifest.description or ""
             self._metadata = copy.deepcopy(self._model_version.manifest.metadata)
-            self._model_schema = copy.deepcopy(
-                self._model_version.manifest.model_schema
-            )
             self._environment = copy.deepcopy(self._model_version.manifest.environment)
             self._framework = (
                 copy.deepcopy(self._model_version.manifest.framework.actual_instance)
@@ -158,7 +157,6 @@ class ModelVersion:
         else:
             self._description = self._model_version.description or ""
             self._metadata = copy.deepcopy(self._model_version.artifact_metadata)
-            self._model_schema = None
             self._environment = None
             self._framework = _ModelFramework.to_model_framework_type(
                 self._model_version.model_framework
@@ -232,23 +230,6 @@ class ModelVersion:
         _validate_artifact_metadata(value)
         self._metadata = copy.deepcopy(value)
 
-    @property
-    def model_schema(self) -> Optional[Dict[str, Any]]:
-        """Get model_schema for the current model"""
-        return self._model_schema
-
-    @model_schema.setter
-    def model_schema(self, value: Optional[Dict[str, Any]]):
-        """set the model_schema for current model"""
-        if not self._model_version.manifest:
-            warnings.warn(
-                message="This model version was created using an older serialization format. model_schema will not be updated",
-                category=DeprecationWarning,
-                stacklevel=2,
-            )
-            return
-        self._model_schema = copy.deepcopy(value)
-
     @property
     def environment(self) -> Optional[ModelVersionEnvironment]:
         """Get the environment details for the model"""
@@ -463,7 +444,6 @@ class ModelVersion:
         if self._model_version.manifest:
             self._model_version.manifest.description = self.description
             self._model_version.manifest.metadata = self.metadata
-            self._model_version.manifest.model_schema = self.model_schema
             self._model_version.manifest.environment = self.environment
             self._model_version.manifest.framework = (
                 Framework.from_dict(self.framework.dict()) if self.framework else None
@@ -498,7 +478,6 @@ def _log_model_version(  # noqa: C901
     progress: Optional[bool] = None,
     framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
     environment: Optional[ModelVersionEnvironment] = None,
-    model_schema: Optional[Dict[str, Any]] = None,
 ) -> ModelVersion:
     if (run and mlfoundry_artifacts_api) or (not run and not mlfoundry_artifacts_api):
         raise MlFoundryException(
@@ -610,7 +589,6 @@ def _log_model_version(  # noqa: C901
         framework=Framework.from_dict(framework.dict()) if framework else None,
         environment=environment,
         step=step,
-        model_schema=model_schema,
     )
     artifact_version_response = mlfoundry_artifacts_api.finalize_artifact_version_post(
         finalize_artifact_version_request_dto=FinalizeArtifactVersionRequestDto(
@@ -624,3 +602,18 @@ def _log_model_version(  # noqa: C901
         )
     )
     return ModelVersion.from_fqn(fqn=artifact_version_response.artifact_version.fqn)
+
+
+def infer_signature(
+    model_input: Any = None,
+    model_output: Optional[
+        Union["pd.DataFrame", "np.ndarray", Dict[str, "np.ndarray"]]
+    ] = 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/ml/mlfoundry_api.py

@@ -1265,7 +1265,6 @@ class MlFoundry:
         progress: Optional[bool] = None,
         framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
         environment: Optional[ModelVersionEnvironment] = None,
-        model_schema: Optional[Dict[str, Any]] = None,
     ) -> ModelVersion:
         """
         Serialize and log a versioned model under the current ml_repo. Each logged model generates a new version
@@ -1390,7 +1389,6 @@ class MlFoundry:
             progress=progress,
             framework=framework,
             environment=environment,
-            model_schema=model_schema,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version

truefoundry/ml/mlfoundry_run.py

@@ -934,7 +934,6 @@ class MlFoundryRun:
         progress: Optional[bool] = None,
         framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
         environment: Optional[ModelVersionEnvironment] = None,
-        model_schema: Optional[Dict[str, Any]] = None,
     ) -> ModelVersion:
         # TODO (chiragjn): Document mapping of framework to list of valid model save kwargs
         # TODO (chiragjn): Add more examples
@@ -1061,7 +1060,6 @@ class MlFoundryRun:
             progress=progress,
             framework=framework,
             environment=environment,
-            model_schema=model_schema,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version

{truefoundry-0.5.1rc1.dist-info → truefoundry-0.5.1rc3.dist-info}/METADATA

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

{truefoundry-0.5.1rc1.dist-info → truefoundry-0.5.1rc3.dist-info}/RECORD

@@ -107,11 +107,11 @@ truefoundry/deploy/v2/lib/models.py,sha256=pSolLMTArDuYpeNsmeeS5DWliloN_iCDfZSpR
 truefoundry/deploy/v2/lib/patched_models.py,sha256=NTU8J_CwdvEuF9zNXwFyN3suOp_196Wrm75DDy5qcXE,14184
 truefoundry/deploy/v2/lib/source.py,sha256=VHKuFREiixUP40D3Mrz-TA70spu1M0RbCzl--qwlFaY,9263
 truefoundry/logger.py,sha256=u-YCNjg5HBwE70uQcpjIG64Ghos-K2ulTWaxC03BSj4,714
-truefoundry/ml/__init__.py,sha256=qfYk5FVNDl50_RJBvIjfzQakOQ4HHLH8oVYmb4SOOh0,1963
+truefoundry/ml/__init__.py,sha256=iTuqq6FuVozL0Vyz1B2Xb5PL2LNnXyALXM0S3ehQicQ,2039
 truefoundry/ml/artifact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=FksxhUpRHb9pgWZmAB16DhXqkAL6UIAPA1S3RJUApQU,46201
 truefoundry/ml/autogen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/ml/autogen/client/__init__.py,sha256=xQKSEKW6Kn-F7GRHaDKjKNjle8JTOFiMs4Liq2kh8Rc,19383
+truefoundry/ml/autogen/client/__init__.py,sha256=uf2QeBvpQi38iGSPkAmSRki2b8nJIBx-3AnJOo9_mHM,19571
 truefoundry/ml/autogen/client/api/__init__.py,sha256=NyMBxBmIzB1o5LzZZwz9LiydHB3-hPqc_sevsnY6Jrw,746
 truefoundry/ml/autogen/client/api/auth_api.py,sha256=zpWzJhUmW6HHMY_atlUf0B25k77E1kue2hmix5I5Ih0,7017
 truefoundry/ml/autogen/client/api/deprecated_api.py,sha256=mu5x_skNcwz8v1Tr6VP72-tVP7pmBV9muyKy_2NH3F0,38824
@@ -126,7 +126,7 @@ truefoundry/ml/autogen/client/api_client.py,sha256=M31IycWorn10FYS8WbO_Whaipr0tV
 truefoundry/ml/autogen/client/api_response.py,sha256=KRyvecPMXF05PaxILHZ8JHoP4rgKBjKONMgG83aU-rM,844
 truefoundry/ml/autogen/client/configuration.py,sha256=V1oaEnxt-NfpaNmp-EZpf2glovzVhM2coWYt8HBNB4M,15723
 truefoundry/ml/autogen/client/exceptions.py,sha256=XbCbDHhYT3BVejdoGNPgEa4oS56ypkwFdxk1iOc_tFY,5355
-truefoundry/ml/autogen/client/models/__init__.py,sha256=_nSy0foDUUFyyE5zSBYKz16d5BvbWzJfCCGvMraUieY,18005
+truefoundry/ml/autogen/client/models/__init__.py,sha256=iX7kGG97Ly93SQfSwFFbggcXAf2OnO9ubNPl0OhN7W4,18193
 truefoundry/ml/autogen/client/models/add_custom_metrics_to_model_version_request_dto.py,sha256=_ISDspicTGjBCYYXubKfRYYSSQVyW3AvG-jFh47-Zfc,2163
 truefoundry/ml/autogen/client/models/add_features_to_model_version_request_dto.py,sha256=OT1-98DyWNfAHz_EgD2gMX2XkrGQ4Re945fhoAl8qSE,2099
 truefoundry/ml/autogen/client/models/agent.py,sha256=lUAbs092yo9hZmqzX6atKebDNf1uEwTf5jLvjOdvoeM,3872
@@ -206,7 +206,7 @@ truefoundry/ml/autogen/client/models/h2_o_framework.py,sha256=5X-0Xbt4tpecV_q6Ww
 truefoundry/ml/autogen/client/models/http_validation_error.py,sha256=bU-uUlz_qeFba3iHdBT7cHjLu6WZ-JmImYqgZ-kCvxM,2466
 truefoundry/ml/autogen/client/models/image_content_part.py,sha256=m_LE5XqD_UD0f8jtWt0adRgz_FoLEvgZbCDTSYORnCc,2717
 truefoundry/ml/autogen/client/models/image_url.py,sha256=TN8iEdBUmtKdTFlApUodaC4o6X_Ggd3OSOKVPv7laJE,2170
-truefoundry/ml/autogen/client/models/infer_method_name.py,sha256=01flpfaOwmqXYFZnrSJCVNl8XK4Sen5slxA6zCo4G18,788
+truefoundry/ml/autogen/client/models/infer_method_name.py,sha256=Cpg3aRGrZCG_6JypR0RMvH0cUqRvpMXxr9rabUGxFuY,788
 truefoundry/ml/autogen/client/models/internal_metadata.py,sha256=6OizrIzA3QxlEnwNFXCZGOMJuUDSYHctl_30q27uBDI,6232
 truefoundry/ml/autogen/client/models/keras_framework.py,sha256=sBKDT4Dzkvr3Rf7c8YKRbtzaO_KrPNsW8tdWI_n701Q,2150
 truefoundry/ml/autogen/client/models/latest_run_log_dto.py,sha256=wE8T8bANTb9u14Jv7DNyKWekZiUAvrzvTcE_H1cRhn4,2326
@@ -247,7 +247,7 @@ truefoundry/ml/autogen/client/models/model_response_dto.py,sha256=osrTxfygkuhxWj
 truefoundry/ml/autogen/client/models/model_server.py,sha256=N5mcWmnX0KzTMtTH8A9h_FmiQcMld3B8CXD_UYeDKaA,700
 truefoundry/ml/autogen/client/models/model_version_dto.py,sha256=9Mh4d8V33gEGYr4sys5lYYDA3bHKajR3oAZLZMAy_io,7094
 truefoundry/ml/autogen/client/models/model_version_environment.py,sha256=wTayQjMen24eJCqXQugzoXLANY5NXdVsAvZSh-QIhJA,2827
-truefoundry/ml/autogen/client/models/model_version_manifest.py,sha256=1FZ9VuMp4THs9bOrVM7YWs4eDHfQTcVTQZSN5UbJClI,4906
+truefoundry/ml/autogen/client/models/model_version_manifest.py,sha256=gRsrZzfLm1cHk5cbHmYQtShXGLLxPKcxRMPDWCqWTZc,4657
 truefoundry/ml/autogen/client/models/model_version_response_dto.py,sha256=D6XOCyggxqTkbePuypqYSHYh1PYeDv7R_J32q61sDM0,2320
 truefoundry/ml/autogen/client/models/multi_part_upload_dto.py,sha256=Ckq405vud8RQmMyKCJQJBlW5iXO7Y2mlgo8eVkiMvxg,3738
 truefoundry/ml/autogen/client/models/multi_part_upload_response_dto.py,sha256=VjH0kvl7rMjgDHjYGHnsh7KsZ5-qn-k3ksdGLJ49nIM,2431
@@ -274,7 +274,8 @@ truefoundry/ml/autogen/client/models/serialization_format.py,sha256=CYq9DR1yipEG
 truefoundry/ml/autogen/client/models/set_experiment_tag_request_dto.py,sha256=nrmi_NxLD1fI2gwlpdqFSMnBS11gRkjS4_GQFHgBcXs,2118
 truefoundry/ml/autogen/client/models/set_tag_request_dto.py,sha256=IRgAdMcWBxmjNV6nZJej4pcNfLmZwrelEZ3otwt7eeE,2144
 truefoundry/ml/autogen/client/models/signed_url_dto.py,sha256=9oHoXBj07xTdc04rqOqJO3eOjQWXCyWPhfHg-6qX60w,1897
-truefoundry/ml/autogen/client/models/sklearn_framework.py,sha256=sO3mQRTGKozmkJl4BrukZpNN9bC1WeEEtuxyio4qwn0,3290
+truefoundry/ml/autogen/client/models/sklearn_framework.py,sha256=eWpv-CNyxsD9pIBKga568DpQvCn5Nn_GvNMzY1ElRW4,3424
+truefoundry/ml/autogen/client/models/sklearn_model_schema.py,sha256=RJIbkzvIDxkFTyygQLvUfeg7sCrkhwm9Euo7FaNnrLw,2526
 truefoundry/ml/autogen/client/models/source.py,sha256=ZG2-3oIs_M8mZNPjoIRPzhHqlm63vLCVQC9idUxQ0XU,5180
 truefoundry/ml/autogen/client/models/source1.py,sha256=io35b9uVzqSvOHw6RlCXxM30thN-vHm1upaSiN2i0II,5183
 truefoundry/ml/autogen/client/models/spa_cy_framework.py,sha256=Kyrvhdalnl1uQLoPMGMry2PuFGCcQ8AGp4Cx4B9AGXQ,2150
@@ -301,10 +302,11 @@ truefoundry/ml/autogen/client/models/url.py,sha256=zMyOmdVkp1ANnQnc9GrHt42xlVwES
 truefoundry/ml/autogen/client/models/user_message.py,sha256=3TEM8qH_zT3dgM197bl44uV7m20IWn6sWeLGfStsN90,2784
 truefoundry/ml/autogen/client/models/validation_error.py,sha256=mFjwoc8g2-Usu1HXZhOQKQ4TGvLy4lwCzk8dHrJ69aA,2597
 truefoundry/ml/autogen/client/models/validation_error_loc_inner.py,sha256=nThJ5Gmy8W2Wok-ZOI4sK7uRe1BAkLS0qzq-XZbq8zs,4915
-truefoundry/ml/autogen/client/models/xg_boost_framework.py,sha256=_YYI6stH0_FRkB7zPfS_85hM_9biwYLr2pLhMBFjARA,2892
+truefoundry/ml/autogen/client/models/xg_boost_framework.py,sha256=ZTvdmfX5eB_ZbA7Yy62NYAzbYKOQkPtA-zF5xxRnVfE,3429
+truefoundry/ml/autogen/client/models/xg_boost_model_schema.py,sha256=5bzi5NEeR9HetiHFJwwkmeC4LE2WvGF9hp9dTLPfEJ4,2526
 truefoundry/ml/autogen/client/rest.py,sha256=9goba8qHjQuVx5O_yRaTKu7PvBnb7r7swfy3dwuTEgk,14281
-truefoundry/ml/autogen/client_README.md,sha256=9z0MWex6ROr_RjjQ_yOZj9iXWle0pMvNITY1sR_D6TM,36112
-truefoundry/ml/autogen/entities/artifacts.py,sha256=Gz7lmWcQAKNsfLr3Qb5S0-orXuzLNbX66qMEscpyEwY,19483
+truefoundry/ml/autogen/client_README.md,sha256=8M-8txeNhpwAiXx7z5UEjPYSq2SoqA1EqtCkkKZDCMc,36276
+truefoundry/ml/autogen/entities/artifacts.py,sha256=x3DxHXPS5OpNb1vmuy6p44g2hymq6BXnivFqv6mQ-ds,20113
 truefoundry/ml/autogen/models/__init__.py,sha256=--TGRea9SQBWFfwtcl3ekb1XGfFTdEkQGSG8a2SJ60I,187
 truefoundry/ml/autogen/models/exceptions.py,sha256=q3n7FGBrg_hUy1UyoefhMnhcXUAaqXlLIGHoOVzn_d8,1390
 truefoundry/ml/autogen/models/schema.py,sha256=IhpO9qbygLqEamP3NIt3m90SseJXCOm1ZTqNbNbW-M0,55772
@@ -331,8 +333,7 @@ truefoundry/ml/log_types/artifacts/artifact.py,sha256=6X4lO23DwM7yTgUOY4RRoYMrf9
 truefoundry/ml/log_types/artifacts/constants.py,sha256=qKxQ5mMvJE4j83BvGW3qNTKunxCiBg_EEjTdgbgJtyE,1036
 truefoundry/ml/log_types/artifacts/dataset.py,sha256=a4dxd2EN8p7Ci-cLGGiDOboN3t0395_XhWE1dmTw1Q4,13112
 truefoundry/ml/log_types/artifacts/general_artifact.py,sha256=B4XErLr-m6RmQWtxMTu3wlFRFcqSwPYp6J0OL4Ng6L0,3179
-truefoundry/ml/log_types/artifacts/model.py,sha256=QJqqgDbflTfsmdsNuV30HvRb0fnaP2UaTrLLqL7jXGM,23312
-truefoundry/ml/log_types/artifacts/model_extras.py,sha256=TIE73bLKfwIVzNiVcjmaZ841A70BHBwu4XAM6ZAQRFI,1045
+truefoundry/ml/log_types/artifacts/model.py,sha256=d2Q-61yXpX1mjXX6t2sL0_xKz-DPA3Z8xYDxzMfUSiA,22820
 truefoundry/ml/log_types/artifacts/utils.py,sha256=yhIZYe3Y-J0rG0Wcwql_Qjxe9hlefFbXUfRFqg4E9cM,7502
 truefoundry/ml/log_types/image/__init__.py,sha256=fcOq8yQnNj1rkLcPeIjLXBpdA1WIeiPsXOlAAvMxx7M,76
 truefoundry/ml/log_types/image/constants.py,sha256=wLtGEOA4T5fZHSlOXPuNDLX3lpbCtwlvGKPFk_1fah0,255
@@ -343,8 +344,8 @@ truefoundry/ml/log_types/plot.py,sha256=HuYvvRA5r8V0xAIuuqMME2IHb9d3SfGHUiuEkOP3
 truefoundry/ml/log_types/pydantic_base.py,sha256=eBlw_AEyAz4iJKDP4zgJOCFWcldwQqpf7FADW1jzIQY,272
 truefoundry/ml/log_types/utils.py,sha256=xjJ21jdPScvFmw3TbVh5NCzbzJwaqiXJyiiT4xxX1EI,335
 truefoundry/ml/logger.py,sha256=VT-BF3BnBYTWVq87O58F0c8uXMu94gYzsiFlGY3_7Ao,458
-truefoundry/ml/mlfoundry_api.py,sha256=HbtMh8uFFueJmxcmYTohV3GnziqbSGkOaq3Nv_I-1eU,62458
-truefoundry/ml/mlfoundry_run.py,sha256=C4fhMc9o1m_ubkRsR8R8rk888ktpqWd0pkqZZQ-Pr4k,44691
+truefoundry/ml/mlfoundry_api.py,sha256=PefznXwBPBv8TD4RYid9aiE8WWEkZQAHoaAbvklFoRs,62364
+truefoundry/ml/mlfoundry_run.py,sha256=RLTZGz3htyxNmG1Xat9WevawEIjDoUCbWrSNql8WynI,44597
 truefoundry/ml/model_framework.py,sha256=2JTQ8m5oqy2FaIBqIEgPeDEfwNBfcQHTw-gdcMXF1pA,15033
 truefoundry/ml/run_utils.py,sha256=0W208wSLUrbdfk2pjNcZlkUi9bNxG2JORqoe-5rVqHI,2423
 truefoundry/ml/session.py,sha256=F83GTC5WwGBjnJ69Ct8MqMnlutYc56JCc6YhEY1Wl-A,5394
@@ -366,7 +367,7 @@ truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py,sha256=5mBCIc-ON
 truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py,sha256=Hf6Dk6Fu6P7DqsK5ULgraf9DStjgigf-kjaRAMBW-RU,8680
 truefoundry/workflow/task.py,sha256=ToitYiKcNzFCtOVQwz1W8sRjbR97eVS7vQBdbgUQtKg,1779
 truefoundry/workflow/workflow.py,sha256=WaTqUjhwfAXDWu4E5ehuwAxrCbDJkoAf1oWmR2E9Qy0,4575
-truefoundry-0.5.1rc1.dist-info/METADATA,sha256=nz3SDzyjyfTrEYBTeMRUAUAvOrtcZewEmKq-AKZwp58,2887
-truefoundry-0.5.1rc1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-truefoundry-0.5.1rc1.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
-truefoundry-0.5.1rc1.dist-info/RECORD,,
+truefoundry-0.5.1rc3.dist-info/METADATA,sha256=isGJJmi0DSlh57OQh5ukDKprSYhHxMt59RG1vsjSbhA,2887
+truefoundry-0.5.1rc3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+truefoundry-0.5.1rc3.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
+truefoundry-0.5.1rc3.dist-info/RECORD,,

truefoundry/ml/log_types/artifacts/model_extras.py

@@ -1,48 +0,0 @@
-import enum
-from typing import List, Type, TypeVar
-
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    ModelSchemaDto,
-)
-from truefoundry.ml.exceptions import MlFoundryException
-from truefoundry.pydantic_v1 import BaseModel
-
-T = TypeVar("T")
-
-
-class BaseEnum(enum.Enum):
-    @classmethod
-    def values(cls: Type[T]) -> List[T]:
-        return [member.value for member in cls]
-
-    @classmethod
-    def _missing_(cls: Type[T], value: object):
-        raise MlFoundryException(
-            f"Unknown value for type {cls.__name__}: {value}", status_code=400
-        )
-
-
-@enum.unique
-class CustomMetricValueType(str, BaseEnum):
-    FLOAT = "float"
-
-
-@enum.unique
-class CustomMetricType(str, BaseEnum):
-    METRIC = "metric"
-    PROJECTION = "projection"
-
-
-class CustomMetric(BaseModel):
-    class Config:
-        validate_assignment = True
-        use_enum_values = True
-        extra = "allow"
-
-    name: str
-    value_type: CustomMetricValueType
-    type: CustomMetricType
-
-
-class ModelSchema(ModelSchemaDto):
-    pass