truefoundry 0.4.6__py3-none-any.whl → 0.5.0rc2__py3-none-any.whl

truefoundry/ml/log_types/artifacts/model.py

@@ -5,39 +5,38 @@ import logging
 import os.path
 import tempfile
 import uuid
+import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple, Union
 
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
     MlFoundryArtifactsRepository,
 )
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    AddCustomMetricsToModelVersionRequestDto,
     ArtifactType,
+    ArtifactVersionSerializationFormat,
     CreateArtifactVersionRequestDto,
-    CreateModelVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
+    ExternalModelSource,
     FinalizeArtifactVersionRequestDto,
-    MetricDto,
+    Framework,
     MlfoundryArtifactsApi,
     ModelDto,
     ModelVersionDto,
+    ModelVersionManifest,
     NotifyArtifactVersionFailureDto,
+    TruefoundryModelSource,
     UpdateModelVersionRequestDto,
 )
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    InternalMetadata as InternalMetadataDto,
+from truefoundry.ml.autogen.client import (
+    Source as ModelVersionSource,
 )
 from truefoundry.ml.enums import ModelFramework
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.constants import (
-    FILES_DIR,
     INTERNAL_METADATA_PATH,
-    MODEL_DIR_NAME,
-    MODEL_SCHEMA_UPDATE_FAILURE_HELP,
 )
-from truefoundry.ml.log_types.artifacts.model_extras import CustomMetric, ModelSchema
 from truefoundry.ml.log_types.artifacts.utils import (
     _copy_additional_files,
     _get_src_dest_pairs,
@@ -45,19 +44,18 @@ from truefoundry.ml.log_types.artifacts.utils import (
     _validate_description,
     calculate_total_size,
 )
+from truefoundry.ml.model_framework import ModelFrameworkType, _ModelFramework
 from truefoundry.ml.session import _get_api_client
-from truefoundry.pydantic_v1 import BaseModel, Extra
-from truefoundry.version import __version__
+from truefoundry.pydantic_v1 import BaseModel, Extra, StrictStr
 
 if TYPE_CHECKING:
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
 
+
 logger = logging.getLogger(__name__)
 
 
 # TODO: Support async download and upload
-
-
 class ModelVersionInternalMetadata(BaseModel):
     class Config:
         extra = Extra.allow
@@ -77,6 +75,10 @@ class ModelVersionInternalMetadata(BaseModel):
         return dct
 
 
+class BlobStorageModelDirectory(BaseModel):
+    uri: StrictStr
+
+
 class ModelVersionDownloadInfo(BaseModel):
     download_dir: str
     model_dir: str
@@ -99,8 +101,7 @@ class ModelVersion:
         self._deleted = False
         self._description: str = ""
         self._metadata: Dict[str, Any] = {}
-        self._metrics: List[MetricDto] = []
-        self._set_metrics_attr()
+        self._framework: Optional["Framework"] = None
         self._set_mutable_attrs()
 
     @classmethod
@@ -138,21 +139,20 @@ class ModelVersion:
                 "Model Version was deleted, cannot perform updates on a deleted version"
             )
 
-    def _set_metrics_attr(self):
-        self._metrics = sorted(
-            self._model_version.metrics or [], key=lambda m: m.timestamp
-        )
-
     def _set_mutable_attrs(self):
         self._description = self._model_version.description or ""
         self._metadata = copy.deepcopy(self._model_version.artifact_metadata)
+        self._framework = (
+            copy.deepcopy(self._model_version.manifest.framework)
+            if self._model_version.manifest
+            else None
+        )
 
     def _refetch_model_version(self):
         _model_version = self._mlfoundry_artifacts_api.get_model_version_get(
             id=self._model_version.id
         )
         self._model_version = _model_version.model_version
-        self._set_metrics_attr()
         self._set_mutable_attrs()
 
     def __repr__(self):
@@ -201,6 +201,8 @@ class ModelVersion:
         """set the description of the model"""
         _validate_description(value)
         self._description = value
+        if self._model_version.manifest:
+            self._model_version.manifest.description = value
 
     @property
     def metadata(self) -> Dict[str, Any]:
@@ -212,14 +214,43 @@ class ModelVersion:
         """set the metadata for current model"""
         _validate_artifact_metadata(value)
         self._metadata = value
+        if self._model_version.manifest:
+            self._model_version.manifest.metadata = value
 
     @property
-    def metrics(self) -> Dict[str, Union[float, int]]:
-        """get the metrics for the current version of the model"""
-        metrics_as_kv: Dict[str, Union[float, int]] = {}
-        for metric in self._metrics:
-            metrics_as_kv[metric.key] = metric.value
-        return metrics_as_kv
+    def framework(self) -> Optional["ModelFrameworkType"]:
+        """Get the framework of the model"""
+        return (
+            _ModelFramework.from_dict(self._framework.actual_instance.dict())
+            if self._framework and self._framework.actual_instance
+            else None
+        )
+
+    @framework.setter
+    def framework(
+        self, value: Optional[Union[str, ModelFramework, "ModelFrameworkType"]]
+    ):
+        """set the framework of the model"""
+        model_framework_type: Optional[ModelFrameworkType] = None
+
+        if value is not None:
+            model_framework_type = _ModelFramework.to_model_framework_type(value)
+
+        if (
+            self._model_version.serialization_format
+            == ArtifactVersionSerializationFormat.V2
+        ):
+            self._model_version.manifest.framework = self._framework = (
+                Framework.from_dict(model_framework_type.dict())
+                if model_framework_type
+                else None
+            )
+        else:
+            warnings.warn(
+                message="This model version was created using an older serialization format. Framework will not be updated",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
 
     @property
     def created_at(self) -> datetime.datetime:
@@ -274,11 +305,29 @@ class ModelVersion:
         path: Optional[Union[str, Path]],
         overwrite: bool = False,
         progress: Optional[bool] = None,
-    ) -> Tuple[ModelVersionInternalMetadata, ModelVersionDownloadInfo]:
+    ) -> ModelVersionDownloadInfo:
         self._ensure_not_deleted()
         download_dir = self.raw_download(
             path=path, overwrite=overwrite, progress=progress
         )
+
+        if (
+            self._model_version.serialization_format
+            == ArtifactVersionSerializationFormat.V2
+        ):
+            model_framework = (
+                ModelFramework(self.framework.type)
+                if self.framework
+                else ModelFramework.UNKNOWN
+            )
+            download_info = ModelVersionDownloadInfo(
+                download_dir=download_dir,
+                model_dir=download_dir,
+                model_framework=model_framework,
+                model_filename=None,
+            )
+            return download_info
+
         internal_metadata_path = os.path.join(download_dir, INTERNAL_METADATA_PATH)
         if not os.path.exists(internal_metadata_path):
             raise MlFoundryException(
@@ -290,12 +339,14 @@ class ModelVersion:
         download_info = ModelVersionDownloadInfo(
             download_dir=os.path.join(download_dir, internal_metadata.files_dir),
             model_dir=os.path.join(
-                download_dir, internal_metadata.files_dir, internal_metadata.model_dir
+                download_dir,
+                internal_metadata.files_dir,
+                internal_metadata.model_dir,
             ),
             model_framework=internal_metadata.framework,
             model_filename=internal_metadata.model_filename,
         )
-        return internal_metadata, download_info
+        return download_info
 
     def download(
         self,
@@ -331,7 +382,7 @@ class ModelVersion:
             print(download_info.model_dir)
             ```
         """
-        _, download_info = self._download(
+        download_info = self._download(
             path=path, overwrite=overwrite, progress=progress
         )
         return download_info
@@ -383,32 +434,36 @@ class ModelVersion:
                 id=self._model_version.id,
                 description=self.description,
                 artifact_metadata=self.metadata,
+                manifest=self._model_version.manifest,
             )
         )
         self._model_version = _model_version.model_version
-        self._set_metrics_attr()
         self._set_mutable_attrs()
 
 
 def _log_model_version(  # noqa: C901
     run: Optional["MlFoundryRun"],
     name: str,
-    model_file_or_folder: str,
-    framework: Optional[Union[ModelFramework, str]],
+    model_file_or_folder: Union[str, BlobStorageModelDirectory],
     mlfoundry_artifacts_api: Optional[MlfoundryArtifactsApi] = None,
     ml_repo_id: Optional[str] = None,
     additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
     description: Optional[str] = None,
     metadata: Optional[Dict[str, Any]] = None,
-    model_schema: Optional[Union[Dict[str, Any], ModelSchema]] = None,
-    custom_metrics: Optional[List[Union[CustomMetric, Dict[str, Any]]]] = None,
     step: Optional[int] = 0,
     progress: Optional[bool] = None,
+    framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
 ) -> ModelVersion:
     if (run and mlfoundry_artifacts_api) or (not run and not mlfoundry_artifacts_api):
         raise MlFoundryException(
             "Exactly one of run, mlfoundry_artifacts_api should be passed"
         )
+
+    if not isinstance(model_file_or_folder, (str, BlobStorageModelDirectory)):
+        raise MlFoundryException(
+            "model_file_or_folder should be of type str or BlobStorageModelDirectory"
+        )
+
     if mlfoundry_artifacts_api and not ml_repo_id:
         raise MlFoundryException(
             "If mlfoundry_artifacts_api is passed, ml_repo_id must also be passed"
@@ -418,89 +473,41 @@ def _log_model_version(  # noqa: C901
 
     assert mlfoundry_artifacts_api is not None
 
-    custom_metrics = custom_metrics or []
+    step = step or 0
+    total_size = None
     metadata = metadata or {}
     additional_files = additional_files or {}
-    step = step or 0
-
-    # validations
-    if framework is None:
-        framework = ModelFramework.UNKNOWN
-    elif not isinstance(framework, ModelFramework):
-        framework = ModelFramework(framework)
 
     _validate_description(description)
     _validate_artifact_metadata(metadata)
 
-    if model_schema is not None and not isinstance(model_schema, ModelSchema):
-        model_schema = ModelSchema.parse_obj(model_schema)
-
-    if custom_metrics and not model_schema:
-        raise MlFoundryException(
-            "Custom Metrics defined without adding the Model Schema"
-        )
-    custom_metrics = [
-        CustomMetric.parse_obj(cm) if not isinstance(cm, CustomMetric) else cm
-        for cm in custom_metrics
-    ]
-
-    logger.info("Logging model and additional files, this might take a while ...")
-    temp_dir = tempfile.TemporaryDirectory(prefix="truefoundry-")
-
-    internal_metadata = ModelVersionInternalMetadata(
-        framework=framework,
-        files_dir=FILES_DIR,
-        model_dir=MODEL_DIR_NAME,
-        model_filename=(
-            os.path.basename(model_file_or_folder)
-            if model_file_or_folder and os.path.isfile(model_file_or_folder)
-            else None
-        ),
-        mlfoundry_version=__version__,
-        truefoundry_version=__version__,
-    )
-
-    try:
-        local_files_dir = os.path.join(temp_dir.name, internal_metadata.files_dir)
-        os.makedirs(local_files_dir, exist_ok=True)
-        # in case model was None, we still create an empty dir
-        local_model_dir = os.path.join(local_files_dir, internal_metadata.model_dir)
-        os.makedirs(local_model_dir, exist_ok=True)
-
-        logger.info("Adding model file/folder to model version content")
-        _model_file_or_folder: Sequence[Tuple[str, str]] = [
-            (model_file_or_folder, MODEL_DIR_NAME.rstrip(os.sep) + os.sep),
-        ]
-
-        temp_dest_to_src_map = _copy_additional_files(
-            root_dir=temp_dir.name,
-            files_dir=internal_metadata.files_dir,
-            model_dir=internal_metadata.model_dir,
-            additional_files=_model_file_or_folder,
-            ignore_model_dir_dest_conflict=True,
-        )
-
-        # verify additional files and paths, copy additional files
-        if additional_files:
-            logger.info("Adding `additional_files` to model version contents")
+    if isinstance(model_file_or_folder, str):
+        logger.info("Logging model and additional files, this might take a while ...")
+        temp_dir = tempfile.TemporaryDirectory(prefix="truefoundry-")
+        try:
+            logger.info("Adding model file/folder to model version content")
             temp_dest_to_src_map = _copy_additional_files(
                 root_dir=temp_dir.name,
-                files_dir=internal_metadata.files_dir,
-                model_dir=internal_metadata.model_dir,
-                additional_files=additional_files,
-                ignore_model_dir_dest_conflict=False,
-                existing_dest_to_src_map=temp_dest_to_src_map,
+                files_dir="",
+                model_dir=None,
+                additional_files=[(model_file_or_folder, "")],
+                ignore_model_dir_dest_conflict=True,
             )
-    except Exception as e:
-        temp_dir.cleanup()
-        raise MlFoundryException("Failed to log model") from e
 
-    # save internal metadata
-    local_internal_metadata_path = os.path.join(temp_dir.name, INTERNAL_METADATA_PATH)
-    os.makedirs(os.path.dirname(local_internal_metadata_path), exist_ok=True)
-    with open(local_internal_metadata_path, "w") as f:
-        json.dump(internal_metadata.dict(), f)
-    temp_dest_to_src_map[local_internal_metadata_path] = local_internal_metadata_path
+            # verify additional files and paths, copy additional files
+            if additional_files:
+                logger.info("Adding `additional_files` to model version contents")
+                temp_dest_to_src_map = _copy_additional_files(
+                    root_dir=temp_dir.name,
+                    files_dir="",
+                    model_dir=None,
+                    additional_files=additional_files,
+                    ignore_model_dir_dest_conflict=False,
+                    existing_dest_to_src_map=temp_dest_to_src_map,
+                )
+        except Exception as e:
+            temp_dir.cleanup()
+            raise MlFoundryException("Failed to log model") from e
 
     # create entry
     _create_artifact_version_response = (
@@ -513,82 +520,65 @@ def _log_model_version(  # noqa: C901
         )
     )
     version_id = _create_artifact_version_response.id
-    artifacts_repo = MlFoundryArtifactsRepository(
-        artifact_identifier=ArtifactIdentifier(
-            artifact_version_id=uuid.UUID(version_id)
-        ),
-        api_client=mlfoundry_artifacts_api.api_client,
-    )
-
-    total_size = calculate_total_size(list(temp_dest_to_src_map.values()))
-    try:
-        logger.info(
-            "Packaging and uploading files to remote with size: %.6f MB",
-            total_size / 1000000.0,
+    artifact_storage_root = _create_artifact_version_response.artifact_storage_root
+    if isinstance(model_file_or_folder, str):
+        # Source is of type TruefoundryModelSource
+        source = ModelVersionSource.from_json(
+            TruefoundryModelSource(
+                type="truefoundry", uri=artifact_storage_root
+            ).to_json()
         )
-        src_dest_pairs = _get_src_dest_pairs(
-            root_dir=temp_dir.name, dest_to_src_map=temp_dest_to_src_map
+        artifacts_repo = MlFoundryArtifactsRepository(
+            artifact_identifier=ArtifactIdentifier(
+                artifact_version_id=uuid.UUID(version_id)
+            ),
+            api_client=mlfoundry_artifacts_api.api_client,
         )
-        artifacts_repo.log_artifacts(src_dest_pairs=src_dest_pairs, progress=progress)
-    except Exception as e:
-        mlfoundry_artifacts_api.notify_failure_post(
-            notify_artifact_version_failure_dto=NotifyArtifactVersionFailureDto(
-                id=version_id
+        total_size = calculate_total_size(list(temp_dest_to_src_map.values()))
+        try:
+            logger.info(
+                "Packaging and uploading files to remote with Total Size: %.6f MB",
+                total_size / 1000000.0,
+            )
+            src_dest_pairs = _get_src_dest_pairs(
+                root_dir=temp_dir.name, dest_to_src_map=temp_dest_to_src_map
+            )
+            artifacts_repo.log_artifacts(
+                src_dest_pairs=src_dest_pairs, progress=progress
             )
+        except Exception as e:
+            mlfoundry_artifacts_api.notify_failure_post(
+                notify_artifact_version_failure_dto=NotifyArtifactVersionFailureDto(
+                    id=version_id
+                )
+            )
+            raise MlFoundryException("Failed to log model") from e
+        finally:
+            temp_dir.cleanup()
+    elif isinstance(model_file_or_folder, BlobStorageModelDirectory):
+        source = ModelVersionSource.from_json(
+            ExternalModelSource(type="external", uri=model_file_or_folder.uri).to_json()
         )
-        raise MlFoundryException("Failed to log model") from e
-    finally:
-        temp_dir.cleanup()
-
-    # Note: Here we call from_dict instead of directly passing in init and relying on it
-    # to convert because the complicated union of types generates a custom type to handle casting
-    # Check the source of `InternalMetadataDto` to see the generated code
-    internal_metadata_dto = InternalMetadataDto.from_dict(
-        internal_metadata.dict() if internal_metadata is not None else {}
+    else:
+        raise MlFoundryException("Invalid model_file_or_folder provided")
+
+    _framework = _ModelFramework.to_model_framework_type(framework)
+    model_manifest = ModelVersionManifest(
+        description=description,
+        metadata=metadata,
+        source=source,
+        framework=Framework.from_dict(_framework.dict()) if _framework else None,
+        step=step if run else 0,
     )
-    mlfoundry_artifacts_api.finalize_artifact_version_post(
+    artifact_version_response = mlfoundry_artifacts_api.finalize_artifact_version_post(
         finalize_artifact_version_request_dto=FinalizeArtifactVersionRequestDto(
             id=version_id,
             run_uuid=run.run_id if run else None,
             artifact_size=total_size,
-            internal_metadata=internal_metadata_dto,
-            step=step if run else None,
-        )
-    )
-    _model_version = mlfoundry_artifacts_api.create_model_version_post(
-        create_model_version_request_dto=CreateModelVersionRequestDto(
-            artifact_version_id=version_id,
-            description=description,
             artifact_metadata=metadata,
-            internal_metadata=internal_metadata_dto,
-            data_path=INTERNAL_METADATA_PATH,
-            step=step if run else None,
+            internal_metadata=None,
+            step=model_manifest.step,
+            manifest=model_manifest,
         )
     )
-    model_version = _model_version.model_version
-
-    # update model schema at end
-    update_args: Dict[str, Any] = {
-        "id": version_id,
-        "model_framework": framework.value,
-    }
-    if model_schema:
-        update_args["model_schema"] = model_schema
-
-    try:
-        _model_version = mlfoundry_artifacts_api.update_model_version_post(
-            update_model_version_request_dto=UpdateModelVersionRequestDto(**update_args)
-        )
-        model_version = _model_version.model_version
-        if model_schema:
-            _model_version = mlfoundry_artifacts_api.add_custom_metrics_to_model_version_post(
-                add_custom_metrics_to_model_version_request_dto=AddCustomMetricsToModelVersionRequestDto(
-                    id=version_id, custom_metrics=custom_metrics
-                )
-            )
-            model_version = _model_version.model_version
-    except Exception:
-        # TODO (chiragjn): what is the best exception to catch here?
-        logger.error(MODEL_SCHEMA_UPDATE_FAILURE_HELP.format(fqn=model_version.fqn))
-
-    return ModelVersion.from_fqn(fqn=model_version.fqn)
+    return ModelVersion.from_fqn(fqn=artifact_version_response.artifact_version.fqn)

truefoundry/ml/mlfoundry_api.py

@@ -2,7 +2,17 @@ import os
 import time
 import uuid
 from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import coolname
 import pandas as pd
@@ -42,7 +52,11 @@ from truefoundry.ml.log_types.artifacts.artifact import (
 )
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
-from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
+from truefoundry.ml.log_types.artifacts.model import (
+    BlobStorageModelDirectory,
+    ModelVersion,
+    _log_model_version,
+)
 from truefoundry.ml.logger import logger
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
 from truefoundry.ml.session import (
@@ -57,6 +71,9 @@ from truefoundry.ml.validation_utils import (
     _validate_run_name,
 )
 
+if TYPE_CHECKING:
+    from truefoundry.ml import ModelFrameworkType
+
 _SEARCH_MAX_RESULTS_DEFAULT = 1000
 
 _INTERNAL_ENV_VARS = [
@@ -1209,12 +1226,12 @@ class MlFoundry:
         *,
         ml_repo: str,
         name: str,
-        model_file_or_folder: str,
-        framework: Optional[Union[ModelFramework, str]],
+        model_file_or_folder: Union[str, BlobStorageModelDirectory],
         additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         progress: Optional[bool] = None,
+        framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
     ) -> ModelVersion:
         """
         Serialize and log a versioned model under the current ml_repo. Each logged model generates a new version
@@ -1226,12 +1243,24 @@ class MlFoundry:
             name (str): Name of the model. If a model with this name already exists under the current ML Repo,
                 the logged model will be added as a new version under that `name`. If no models exist with the given
                 `name`, the given model will be logged as version 1.
-            model_file_or_folder (str): Path to either a single file or a folder containing model files. This folder
-                is usually created using serialization methods of libraries or frameworks e.g. `joblib.dump`,
-                `model.save_pretrained(...)`, `torch.save(...)`, `model.save(...)`
-            framework (Union[enums.ModelFramework, str]): Model Framework. Ex:- pytorch, sklearn, tensorflow etc.
-                The full list of supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
-                Can also be `None` when `model` is `None`.
+
+            model_file_or_folder (Union[str, BlobStorageModelDirectory]):
+                str:
+                    Path to either a single file or a folder containing model files.
+                    This folder is typically created using serialization methods from libraries or frameworks,
+                    e.g., `joblib.dump`, `model.save_pretrained(...)`, `torch.save(...)`, or `model.save(...)`.
+                BlobStorageModelDirectory:
+                    uri (str): URI to the model file or folder in a storage integration associated with the specified ML Repo.
+                        The model files or folder must reside within the same storage integration as the specified ML Repo.
+                        Accepted URI formats include `s3://integration-bucket-name/prefix/path/to/model` or `gs://integration-bucket-name/prefix/path/to/model`.
+                        If the URI points to a model in a different storage integration, an error will be raised indicating "Invalid source URI."
+
+            framework (Optional[Union[ModelFramework, ModelFrameworkType]]): Framework used for model serialization.
+                Supported frameworks values (ModelFrameworkType) can be imported from `from truefoundry.ml import *`.
+                Supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
+                Can also be `None` if the framework is not known or not supported.
+                **Deprecated**: Prefer `ModelFrameworkType` over `enums.ModelFramework`.
+
             additional_files (Sequence[Tuple[Union[str, Path], Optional[str]]], optional): A list of pairs
                 of (source path, destination path) to add additional files and folders
                 to the model version contents. The first member of the pair should be a file or directory path
@@ -1245,10 +1274,12 @@ class MlFoundry:
                 You can also add additional files to model/ subdirectory by specifying the destination path as model/
 
                 ```python
+                from truefoundry.ml import TensorFlowFramework
+
                 run.log_model(
                     name="xyz",
                     model_file_or_folder="clf.joblib",
-                    framework="sklearn",
+                    framework=TensorFlowFramework(),
                     additional_files=[("foo.txt", "foo/bar/foo.txt"), ("tokenizer/", "foo/tokenizer/")]
                 )
                 ```
@@ -1281,8 +1312,7 @@ class MlFoundry:
             ### Sklearn
 
             ```python
-            from truefoundry.ml import get_client
-            from truefoundry.ml.enums import ModelFramework
+            from truefoundry.ml import get_client, SklearnFramework
 
             import joblib
             import numpy as np
@@ -1307,7 +1337,7 @@ class MlFoundry:
                 ml_repo="my-classification-project",
                 name="my-sklearn-model",
                 model_file_or_folder="sklearn-pipeline.joblib",
-                framework=ModelFramework.SKLEARN,
+                framework=SklearnFramework(),
                 metadata={"accuracy": 0.99, "f1": 0.80},
                 step=1,  # step number, useful when using iterative algorithms like SGD
             )
@@ -1317,8 +1347,7 @@ class MlFoundry:
             ### Huggingface Transformers
 
             ```python
-            from truefoundry.ml import get_client
-            from truefoundry.ml.enums import ModelFramework
+            from truefoundry.ml import get_client, TransformersFramework, LibraryName
 
             import torch
             from transformers import AutoTokenizer, AutoConfig, pipeline, AutoModelForCausalLM
@@ -1342,7 +1371,7 @@ class MlFoundry:
                 ml_repo="my-llm-project",
                 name="my-transformers-model",
                 model_file_or_folder="my-transformers-model/",
-                framework=ModelFramework.TRANSFORMERS
+                framework=TransformersFramework(library_name=LibraryName.TRANSFORMERS, pipeline_tag='text-generation')
             )
             print(model_version.fqn)
             ```
@@ -1356,12 +1385,12 @@ class MlFoundry:
             ml_repo_id=ml_repo_id,
             name=name,
             model_file_or_folder=model_file_or_folder,
-            framework=framework,
             additional_files=additional_files,
             description=description,
             metadata=metadata,
             step=None,
             progress=progress,
+            framework=framework,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version