genesis-flow 1.0.0__py3-none-any.whl

mlflow/transformers/model_io.py

@@ -0,0 +1,301 @@
+import logging
+import shutil
+
+from mlflow.environment_variables import (
+    MLFLOW_HUGGINGFACE_DISABLE_ACCELERATE_FEATURES,
+    MLFLOW_HUGGINGFACE_MODEL_MAX_SHARD_SIZE,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_STATE
+from mlflow.transformers.flavor_config import FlavorKey, get_peft_base_model, is_peft_model
+
+_logger = logging.getLogger(__name__)
+
+# File/directory names for saved artifacts
+_MODEL_BINARY_FILE_NAME = "model"
+_COMPONENTS_BINARY_DIR_NAME = "components"
+_PROCESSOR_BINARY_DIR_NAME = "processor"
+
+
+def save_pipeline_pretrained_weights(path, pipeline, flavor_conf, processor=None):
+    """
+    Save the binary artifacts of the pipeline to the specified local path.
+
+    Args:
+        path: The local path to save the pipeline
+        pipeline: Transformers pipeline instance
+        flavor_conf: The flavor configuration constructed for the pipeline
+        processor: Optional processor instance to save alongside the pipeline
+    """
+    model = get_peft_base_model(pipeline.model) if is_peft_model(pipeline.model) else pipeline.model
+
+    model.save_pretrained(
+        save_directory=path.joinpath(_MODEL_BINARY_FILE_NAME),
+        max_shard_size=MLFLOW_HUGGINGFACE_MODEL_MAX_SHARD_SIZE.get(),
+    )
+
+    component_dir = path.joinpath(_COMPONENTS_BINARY_DIR_NAME)
+    for name in flavor_conf.get(FlavorKey.COMPONENTS, []):
+        getattr(pipeline, name).save_pretrained(component_dir.joinpath(name))
+
+    if processor:
+        processor.save_pretrained(component_dir.joinpath(_PROCESSOR_BINARY_DIR_NAME))
+
+
+def save_local_checkpoint(path, checkpoint_dir, flavor_conf, processor=None):
+    """
+    Save the local checkpoint of the model and other components to the specified local path.
+
+    Args:
+        path: The local path to save the pipeline
+        checkpoint_dir: The local path to the checkpoint directory
+        flavor_conf: The flavor configuration constructed for the pipeline
+        processor: Optional processor instance to save alongside the pipeline
+    """
+    # Copy files within checkpoint dir to the model path
+    shutil.copytree(checkpoint_dir, path.joinpath(_MODEL_BINARY_FILE_NAME))
+
+    for name in flavor_conf.get(FlavorKey.COMPONENTS, []):
+        # Other pipeline components such as tokenizer may not saved in the checkpoint.
+        # We first try to load the component instance from the checkpoint directory,
+        # if it fails, we load the component from the HuggingFace Hub.
+        try:
+            component = _load_component(flavor_conf, name, local_path=checkpoint_dir)
+        except Exception:
+            repo_id = flavor_conf[FlavorKey.MODEL_NAME]
+            _logger.info(
+                f"The {name} state file is not found ins the local checkpoint directory. MLflow "
+                f"will use the default component state from the base HF repository {repo_id}."
+            )
+            component = _load_component(flavor_conf, name, repo_id=repo_id)
+
+        component.save_pretrained(path.joinpath(_COMPONENTS_BINARY_DIR_NAME, name))
+
+    if processor:
+        processor.save_pretrained(
+            path.joinpath(_COMPONENTS_BINARY_DIR_NAME, _PROCESSOR_BINARY_DIR_NAME)
+        )
+
+
+def load_model_and_components_from_local(path, flavor_conf, accelerate_conf, device=None):
+    """
+    Load the model and components of a Transformer pipeline from the specified local path.
+
+    Args:
+        path: The local path contains MLflow model artifacts
+        flavor_conf: The flavor configuration
+        accelerate_conf: The configuration for the accelerate library
+        device: The device to load the model onto
+    """
+    loaded = {}
+
+    # NB: Path resolution for models that were saved prior to 2.4.1 release when the patching for
+    #     the saved pipeline or component artifacts was handled by duplicate entries for components
+    #     (artifacts/pipeline/* and artifacts/components/*) and pipelines were saved via the
+    #     "artifacts/pipeline/*" path. In order to load the older formats after the change, the
+    #     presence of the new path key is checked.
+    model_path = path.joinpath(flavor_conf.get(FlavorKey.MODEL_BINARY, "pipeline"))
+    loaded[FlavorKey.MODEL] = _load_model(model_path, flavor_conf, accelerate_conf, device)
+
+    components = flavor_conf.get(FlavorKey.COMPONENTS, [])
+    if FlavorKey.PROCESSOR_TYPE in flavor_conf:
+        components.append("processor")
+
+    for component_key in components:
+        component_path = path.joinpath(_COMPONENTS_BINARY_DIR_NAME, component_key)
+        loaded[component_key] = _load_component(
+            flavor_conf, component_key, local_path=component_path
+        )
+
+    return loaded
+
+
+def load_model_and_components_from_huggingface_hub(flavor_conf, accelerate_conf, device=None):
+    """
+    Load the model and components of a Transformer pipeline from HuggingFace Hub.
+
+    Args:
+        flavor_conf: The flavor configuration
+        accelerate_conf: The configuration for the accelerate library
+        device: The device to load the model onto
+    """
+    loaded = {}
+
+    model_repo = flavor_conf[FlavorKey.MODEL_NAME]
+    model_revision = flavor_conf.get(FlavorKey.MODEL_REVISION)
+
+    if not model_revision:
+        raise MlflowException(
+            "The model was saved with 'save_pretrained' set to False, but the commit hash is not "
+            "found in the saved metadata. Loading the model with the different version may cause "
+            "inconsistency issue and security risk.",
+            error_code=INVALID_STATE,
+        )
+
+    loaded[FlavorKey.MODEL] = _load_model(
+        model_repo, flavor_conf, accelerate_conf, device, revision=model_revision
+    )
+
+    components = flavor_conf.get(FlavorKey.COMPONENTS, [])
+    if FlavorKey.PROCESSOR_TYPE in flavor_conf:
+        components.append("processor")
+
+    for name in components:
+        loaded[name] = _load_component(flavor_conf, name)
+
+    return loaded
+
+
+def _load_component(flavor_conf, name, local_path=None, repo_id=None):
+    import transformers
+
+    _COMPONENT_TO_AUTOCLASS_MAP = {
+        FlavorKey.TOKENIZER: transformers.AutoTokenizer,
+        FlavorKey.FEATURE_EXTRACTOR: transformers.AutoFeatureExtractor,
+        FlavorKey.PROCESSOR: transformers.AutoProcessor,
+        FlavorKey.IMAGE_PROCESSOR: transformers.AutoImageProcessor,
+    }
+
+    component_name = flavor_conf[FlavorKey.COMPONENT_TYPE.format(name)]
+    if hasattr(transformers, component_name):
+        cls = getattr(transformers, component_name)
+        trust_remote = False
+    else:
+        if local_path is None:
+            raise MlflowException(
+                f"A custom component `{component_name}` was specified, "
+                "but no local config file was found to retrieve the "
+                "definition. Make sure your model was saved with "
+                "save_pretrained=True."
+            )
+        cls = _COMPONENT_TO_AUTOCLASS_MAP[name]
+        trust_remote = True
+
+    if local_path is not None:
+        # Load component from local file
+        return cls.from_pretrained(str(local_path), trust_remote_code=trust_remote)
+    else:
+        # Load component from HuggingFace Hub
+        repo = repo_id or flavor_conf[FlavorKey.COMPONENT_NAME.format(name)]
+        revision = flavor_conf.get(FlavorKey.COMPONENT_REVISION.format(name))
+        return cls.from_pretrained(repo, revision=revision, trust_remote_code=trust_remote)
+
+
+def _load_class_from_transformers_config(model_name_or_path, revision=None):
+    """
+    This method retrieves the Transformers AutoClass from the transformers config.
+    Using the correct AutoClass allows us to leverage Transformers' model loading
+    machinery, which is necessary for supporting models using custom code.
+    """
+    import transformers
+    from transformers import AutoConfig
+
+    config = AutoConfig.from_pretrained(
+        model_name_or_path,
+        revision=revision,
+        # trust_remote_code is set to True in order to
+        # make sure the config gets loaded as the correct
+        # class. if this is not set for custom models, the
+        # base class will be loaded instead of the custom one.
+        trust_remote_code=True,
+    )
+
+    # the model's class name (e.g. "MPTForCausalLM")
+    # is stored in the `architectures` field. it
+    # seems to usually just have one element.
+    class_name = config.architectures[0]
+
+    # if the class is available in transformers natively,
+    # then we don't need to execute any custom code.
+    if hasattr(transformers, class_name):
+        cls = getattr(transformers, class_name)
+        return cls, False
+    else:
+        # else, we need to fetch the correct AutoClass.
+        # this is defined in the `auto_map` field. there
+        # should only be one AutoClass that maps to the
+        # model's class name.
+        auto_classes = [
+            auto_class
+            for auto_class, module in config.auto_map.items()
+            if module.split(".")[-1] == class_name
+        ]
+
+        if len(auto_classes) == 0:
+            raise MlflowException(f"Couldn't find a loader class for {class_name}")
+
+        auto_class = auto_classes[0]
+        cls = getattr(transformers, auto_class)
+
+        # we will need to trust remote code when loading the model
+        return cls, True
+
+
+def _load_model(model_name_or_path, flavor_conf, accelerate_conf, device, revision=None):
+    """
+    Try to load a model with various loading strategies.
+      1. Try to load the model with accelerate
+      2. Try to load the model with the specified device
+      3. Load the model without the device
+    """
+    import transformers
+
+    if hasattr(transformers, flavor_conf[FlavorKey.MODEL_TYPE]):
+        cls = getattr(transformers, flavor_conf[FlavorKey.MODEL_TYPE])
+        trust_remote = False
+    else:
+        cls, trust_remote = _load_class_from_transformers_config(
+            model_name_or_path, revision=revision
+        )
+
+    load_kwargs = {"revision": revision} if revision else {}
+    if trust_remote:
+        load_kwargs.update({"trust_remote_code": True})
+
+    if model := _try_load_model_with_accelerate(
+        cls, model_name_or_path, {**accelerate_conf, **load_kwargs}
+    ):
+        return model
+
+    load_kwargs["device"] = device
+    if torch_dtype := flavor_conf.get(FlavorKey.TORCH_DTYPE):
+        load_kwargs[FlavorKey.TORCH_DTYPE] = torch_dtype
+
+    if model := _try_load_model_with_device(cls, model_name_or_path, load_kwargs):
+        return model
+    _logger.warning(
+        "Could not specify device parameter for this pipeline type."
+        "Falling back to loading the model with the default device."
+    )
+
+    load_kwargs.pop("device", None)
+    return cls.from_pretrained(model_name_or_path, **load_kwargs)
+
+
+def _try_load_model_with_accelerate(model_class, model_name_or_path, load_kwargs):
+    if MLFLOW_HUGGINGFACE_DISABLE_ACCELERATE_FEATURES.get():
+        return None
+
+    try:
+        return model_class.from_pretrained(model_name_or_path, **load_kwargs)
+    except (ValueError, TypeError, NotImplementedError, ImportError):
+        # NB: ImportError is caught here in the event that `accelerate` is not installed
+        # on the system, which will raise if `low_cpu_mem_usage` is set or the argument
+        # `device_map` is set and accelerate is not installed.
+        pass
+
+
+def _try_load_model_with_device(model_class, model_name_or_path, load_kwargs):
+    try:
+        return model_class.from_pretrained(model_name_or_path, **load_kwargs)
+    except OSError as e:
+        revision = load_kwargs.get("revision")
+        if f"{revision} is not a valid git identifier" in str(e):
+            raise MlflowException(
+                f"The model was saved with a HuggingFace Hub repository name '{model_name_or_path}'"
+                f"and a commit hash '{revision}', but the commit is not found in the repository. "
+            )
+        else:
+            raise e
+    except (ValueError, TypeError, NotImplementedError):
+        pass

mlflow/transformers/peft.py

@@ -0,0 +1,51 @@
+"""
+PEFT (Parameter-Efficient Fine-Tuning) is a library for efficiently adapting large pretrained
+models without fine-tuning all of model parameters but only a small number of (extra) parameters.
+Users can define a PEFT model that wraps a Transformer model to apply a thin adapter layer on
+top of the base model. The PEFT model provides almost the same APIs as the original model such
+as from_pretrained(), save_pretrained().
+"""
+
+_PEFT_ADAPTOR_DIR_NAME = "peft"
+
+
+def is_peft_model(model) -> bool:
+    try:
+        from peft import PeftModel
+    except ImportError:
+        return False
+
+    return isinstance(model, PeftModel)
+
+
+def get_peft_base_model(model):
+    """Extract the base model from a PEFT model."""
+    peft_config = model.peft_config.get(model.active_adapter) if model.peft_config else None
+
+    # PEFT usually wraps the base model with two additional classes, one is PeftModel class
+    # and the other is the adaptor specific class, like LoraModel class, so the class hierarchy
+    # looks like PeftModel -> LoraModel -> BaseModel
+    # However, when the PEFT config is the one for "prompt learning", there is not adaptor class
+    # and the PeftModel class directly wraps the base model.
+    if peft_config and not peft_config.is_prompt_learning:
+        return model.base_model.model
+
+    return model.base_model
+
+
+def get_model_with_peft_adapter(base_model, peft_adapter_path):
+    """
+    Apply the PEFT adapter to the base model to create a PEFT model.
+
+    NB: The alternative way to load PEFT adapter is to use load_adapter API like
+    `base_model.load_adapter(peft_adapter_path)`, as it injects the adapter weights
+    into the model in-place hence reducing the memory footprint. However, doing so
+    returns the base model class and not the PEFT model, loosing some properties
+    such as peft_config. This is not preferable because load_model API should
+    return the exact same object that was saved. Hence we construct the PEFT model
+    instead of in-place injection, for consistency over the memory saving which
+    should be small in most cases.
+    """
+    from peft import PeftModel
+
+    return PeftModel.from_pretrained(base_model, peft_adapter_path)

mlflow/transformers/signature.py

@@ -0,0 +1,183 @@
+import json
+import logging
+
+import numpy as np
+
+from mlflow.environment_variables import MLFLOW_INPUT_EXAMPLE_INFERENCE_TIMEOUT
+from mlflow.models.signature import ModelSignature, infer_signature
+from mlflow.models.utils import _contains_params
+from mlflow.types.schema import ColSpec, DataType, Schema, TensorSpec
+from mlflow.utils.os import is_windows
+from mlflow.utils.timeout import MlflowTimeoutError, run_with_timeout
+
+_logger = logging.getLogger(__name__)
+
+
+_TEXT2TEXT_SIGNATURE = ModelSignature(
+    inputs=Schema([ColSpec("string")]),
+    outputs=Schema([ColSpec("string")]),
+)
+_CLASSIFICATION_SIGNATURE = ModelSignature(
+    inputs=Schema([ColSpec("string")]),
+    outputs=Schema([ColSpec("string", name="label"), ColSpec("double", name="score")]),
+)
+
+# Order is important here, the first matching task type will be used
+_DEFAULT_SIGNATURE_FOR_TASK = {
+    "token-classification": _TEXT2TEXT_SIGNATURE,
+    "translation": _TEXT2TEXT_SIGNATURE,
+    "text-generation": _TEXT2TEXT_SIGNATURE,
+    "text2text-generation": _TEXT2TEXT_SIGNATURE,
+    "text-classification": _CLASSIFICATION_SIGNATURE,
+    "conversational": _TEXT2TEXT_SIGNATURE,
+    "fill-mask": _TEXT2TEXT_SIGNATURE,
+    "summarization": _TEXT2TEXT_SIGNATURE,
+    "image-classification": _CLASSIFICATION_SIGNATURE,
+    "zero-shot-classification": ModelSignature(
+        inputs=Schema(
+            [
+                ColSpec(DataType.string, name="sequences"),
+                ColSpec(DataType.string, name="candidate_labels"),
+                ColSpec(DataType.string, name="hypothesis_template"),
+            ]
+        ),
+        outputs=Schema(
+            [
+                ColSpec(DataType.string, name="sequence"),
+                ColSpec(DataType.string, name="labels"),
+                ColSpec(DataType.double, name="scores"),
+            ]
+        ),
+    ),
+    "automatic-speech-recognition": ModelSignature(
+        inputs=Schema([ColSpec(DataType.binary)]),
+        outputs=Schema([ColSpec(DataType.string)]),
+    ),
+    "audio-classification": ModelSignature(
+        inputs=Schema([ColSpec(DataType.binary)]),
+        outputs=Schema(
+            [ColSpec(DataType.double, name="score"), ColSpec(DataType.string, name="label")]
+        ),
+    ),
+    "table-question-answering": ModelSignature(
+        inputs=Schema(
+            [ColSpec(DataType.string, name="query"), ColSpec(DataType.string, name="table")]
+        ),
+        outputs=Schema([ColSpec(DataType.string)]),
+    ),
+    "question-answering": ModelSignature(
+        inputs=Schema(
+            [ColSpec(DataType.string, name="question"), ColSpec(DataType.string, name="context")]
+        ),
+        outputs=Schema([ColSpec(DataType.string)]),
+    ),
+    "feature-extraction": ModelSignature(
+        inputs=Schema([ColSpec(DataType.string)]),
+        outputs=Schema([TensorSpec(np.dtype("float64"), [-1], "double")]),
+    ),
+}
+
+
+def infer_or_get_default_signature(
+    pipeline, example=None, model_config=None, flavor_config=None
+) -> ModelSignature:
+    """
+    Assigns a default ModelSignature for a given Pipeline type that has pyfunc support. These
+    default signatures should only be generated and assigned when saving a model iff the user
+    has not supplied a signature.
+    For signature inference in some Pipelines that support complex input types, an input example
+    is needed.
+    """
+    import transformers
+
+    if example is not None and isinstance(pipeline, transformers.Pipeline):
+        try:
+            timeout = MLFLOW_INPUT_EXAMPLE_INFERENCE_TIMEOUT.get()
+            if timeout and is_windows():
+                timeout = None
+                _logger.warning(
+                    "On Windows, timeout is not supported for model signature inference. "
+                    "Therefore, the operation is not bound by a timeout and may hang indefinitely. "
+                    "If it hangs, please consider specifying the signature manually."
+                )
+            return _infer_signature_with_example(
+                pipeline, example, model_config, flavor_config, timeout
+            )
+        except Exception as e:
+            if isinstance(e, MlflowTimeoutError):
+                msg = (
+                    "Attempted to generate a signature for the saved pipeline but prediction timed "
+                    f"out after {timeout} seconds. Falling back to the default signature for the "
+                    "pipeline. You can specify a signature manually or increase the timeout "
+                    f"by setting the environment variable {MLFLOW_INPUT_EXAMPLE_INFERENCE_TIMEOUT}"
+                )
+            else:
+                msg = (
+                    "Attempted to generate a signature for the saved pipeline but encountered an "
+                    f"error. Fall back to the default signature for the pipeline type. Error: {e}"
+                )
+            _logger.warning(msg)
+
+    task = getattr(pipeline, "task", None)
+    if task.startswith("translation_"):
+        task = "translation"
+    if signature := _DEFAULT_SIGNATURE_FOR_TASK.get(task):
+        return signature
+
+    _logger.warning(
+        "An unsupported task type was supplied for signature inference. Either provide an "
+        "`input_example` or generate a signature manually via `infer_signature` to have a "
+        "signature recorded in the MLmodel file."
+    )
+
+
+def _infer_signature_with_example(
+    pipeline, example, model_config=None, flavor_config=None, timeout=None
+) -> ModelSignature:
+    params = None
+    if _contains_params(example):
+        example, params = example
+    example = format_input_example_for_special_cases(example, pipeline)
+
+    if timeout:
+        _logger.info(
+            "Running model prediction to infer the model output signature with a timeout "
+            f"of {timeout} seconds. You can specify a different timeout by setting the "
+            f"environment variable {MLFLOW_INPUT_EXAMPLE_INFERENCE_TIMEOUT}."
+        )
+        with run_with_timeout(timeout):
+            prediction = generate_signature_output(
+                pipeline, example, model_config, flavor_config, params
+            )
+    else:
+        prediction = generate_signature_output(
+            pipeline, example, model_config, flavor_config, params
+        )
+    return infer_signature(example, prediction, params)
+
+
+def format_input_example_for_special_cases(input_example, pipeline):
+    """
+    Handles special formatting for specific types of Pipelines so that the displayed example
+    reflects the correct example input structure that mirrors the behavior of the input parsing
+    for pyfunc.
+    """
+    input_data = input_example[0] if isinstance(input_example, tuple) else input_example
+
+    if (
+        pipeline.task == "zero-shot-classification"
+        and isinstance(input_data, dict)
+        and isinstance(input_data["candidate_labels"], list)
+    ):
+        input_data["candidate_labels"] = json.dumps(input_data["candidate_labels"])
+    return input_data if not isinstance(input_example, tuple) else (input_data, input_example[1])
+
+
+def generate_signature_output(pipeline, data, model_config=None, flavor_config=None, params=None):
+    # Lazy import to avoid circular dependencies. Ideally we should move _TransformersWrapper
+    # out from __init__.py to avoid this.
+    from mlflow.transformers import _TransformersWrapper
+
+    return _TransformersWrapper(
+        pipeline=pipeline, model_config=model_config, flavor_config=flavor_config
+    ).predict(data, params=params)

mlflow/transformers/torch_utils.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+if TYPE_CHECKING:
+    import torch
+
+_TORCH_DTYPE_KEY = "torch_dtype"
+
+
+def _extract_torch_dtype_if_set(pipeline) -> Optional[torch.dtype]:
+    """
+    Extract the torch datatype argument if set and return as a string encoded value.
+    """
+    try:
+        import torch
+    except ImportError:
+        # If torch is not installed, safe to assume the model doesn't have a custom torch_dtype
+        return None
+
+    # Check model dtype as pipeline's torch_dtype field doesn't always reflect the model's dtype
+    model_dtype = pipeline.model.dtype if hasattr(pipeline.model, "dtype") else None
+
+    # If the underlying model is PyTorch model, dtype must be a torch.dtype instance
+    return model_dtype if isinstance(model_dtype, torch.dtype) else None
+
+
+def _deserialize_torch_dtype(dtype_str: str) -> torch.dtype:
+    """
+    Convert the string-encoded `torch_dtype` pipeline argument back to the correct `torch.dtype`
+    instance value for applying to a loaded pipeline instance.
+    """
+    try:
+        import torch
+    except ImportError as e:
+        raise MlflowException(
+            "Unable to determine if the value supplied by the argument "
+            "torch_dtype is valid since torch is not installed.",
+            error_code=INVALID_PARAMETER_VALUE,
+        ) from e
+
+    if dtype_str.startswith("torch."):
+        dtype_str = dtype_str[6:]
+
+    dtype = getattr(torch, dtype_str, None)
+    if isinstance(dtype, torch.dtype):
+        return dtype
+
+    raise MlflowException(
+        f"The value '{dtype_str}' is not a valid torch.dtype",
+        error_code=INVALID_PARAMETER_VALUE,
+    )

mlflow/types/__init__.py

@@ -0,0 +1,21 @@
+"""
+The :py:mod:`mlflow.types` module defines data types and utilities to be used by other mlflow
+components to describe interface independent of other frameworks or languages.
+"""
+
+from mlflow.version import IS_TRACING_SDK_ONLY
+
+if not IS_TRACING_SDK_ONLY:
+    import mlflow.types.llm  # noqa: F401
+
+    # Our typing system depends on numpy, which is not included in mlflow-tracing package
+    from mlflow.types.schema import ColSpec, DataType, ParamSchema, ParamSpec, Schema, TensorSpec
+
+    __all__ = [
+        "Schema",
+        "ColSpec",
+        "DataType",
+        "TensorSpec",
+        "ParamSchema",
+        "ParamSpec",
+    ]