chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,519 @@
|
||||
import textwrap
|
||||
import warnings
|
||||
from typing import Any
|
||||
|
||||
from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS
|
||||
from mlflow.utils.autologging_utils.versioning import (
|
||||
get_min_max_version_and_pip_release,
|
||||
)
|
||||
|
||||
|
||||
def _create_placeholder(key: str):
|
||||
return "{{ " + key + " }}"
|
||||
|
||||
|
||||
def _replace_keys_with_placeholders(d: dict[str, Any]) -> dict[str, Any]:
|
||||
return {_create_placeholder(k): v for k, v in d.items()}
|
||||
|
||||
|
||||
def _get_indentation_of_key(line: str, placeholder: str) -> str:
|
||||
index = line.find(placeholder)
|
||||
return (index * " ") if index != -1 else ""
|
||||
|
||||
|
||||
def _indent(text: str, indent: str) -> str:
|
||||
"""Indent everything but first line in text."""
|
||||
lines = text.splitlines()
|
||||
if len(lines) <= 1:
|
||||
return text
|
||||
|
||||
else:
|
||||
first_line = lines[0]
|
||||
subsequent_lines = "\n".join(list(lines[1:]))
|
||||
indented_subsequent_lines = textwrap.indent(subsequent_lines, indent)
|
||||
return first_line + "\n" + indented_subsequent_lines
|
||||
|
||||
|
||||
def _replace_all(text: str, replacements: dict[str, str]) -> str:
|
||||
"""
|
||||
Replace all instances of replacements.keys() with their corresponding
|
||||
values in text. The replacements will be inserted on the same line
|
||||
with wrapping to the same level of indentation, for example:
|
||||
|
||||
```
|
||||
Args:
|
||||
param_1: {{ key }}
|
||||
```
|
||||
|
||||
will become...
|
||||
|
||||
```
|
||||
Args:
|
||||
param_1: replaced_value_at same indentation as prior
|
||||
and if there are more lines they will also
|
||||
have the same indentation.
|
||||
```
|
||||
"""
|
||||
for key, value in replacements.items():
|
||||
if key in text:
|
||||
indent = _get_indentation_of_key(text, key)
|
||||
indented_value = _indent(value, indent)
|
||||
text = text.replace(key, indented_value)
|
||||
return text
|
||||
|
||||
|
||||
class ParamDocs(dict):
|
||||
"""
|
||||
Represents a set of parameter documents in the docstring.
|
||||
"""
|
||||
|
||||
def __repr__(self):
|
||||
return f"ParamDocs({super().__repr__()})"
|
||||
|
||||
def format(self, **kwargs):
|
||||
"""
|
||||
Formats values to be substituted in via the format_docstring() method.
|
||||
|
||||
Args:
|
||||
kwargs: A `dict` in the form of `{"< placeholder name >": "< value >"}`.
|
||||
|
||||
Returns:
|
||||
A new `ParamDocs` instance with the formatted param docs.
|
||||
|
||||
.. code-block:: text
|
||||
:caption: Example
|
||||
|
||||
>>> pd = ParamDocs(p1="{{ doc1 }}", p2="{{ doc2 }}")
|
||||
>>> pd.format(doc1="foo", doc2="bar")
|
||||
ParamDocs({'p1': 'foo', 'p2': 'bar'})
|
||||
"""
|
||||
replacements = _replace_keys_with_placeholders(kwargs)
|
||||
return ParamDocs({k: _replace_all(v, replacements) for k, v in self.items()})
|
||||
|
||||
def format_docstring(self, docstring: str) -> str:
|
||||
"""
|
||||
Formats placeholders in `docstring`.
|
||||
|
||||
Args:
|
||||
docstring: A docstring with placeholders to be replaced.
|
||||
If provided with None, will return None.
|
||||
|
||||
.. code-block:: text
|
||||
:caption: Example
|
||||
|
||||
>>> pd = ParamDocs(p1="doc1", p2="doc2
|
||||
doc2 second line")
|
||||
>>> docstring = '''
|
||||
... Args:
|
||||
... p1: {{ p1 }}
|
||||
... p2: {{ p2 }}
|
||||
... '''.strip()
|
||||
>>> print(pd.format_docstring(docstring))
|
||||
"""
|
||||
if docstring is None:
|
||||
return None
|
||||
|
||||
replacements = _replace_keys_with_placeholders(self)
|
||||
lines = docstring.splitlines()
|
||||
for i, line in enumerate(lines):
|
||||
lines[i] = _replace_all(line, replacements)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def format_docstring(param_docs):
|
||||
"""
|
||||
Returns a decorator that replaces param doc placeholders (e.g. '{{ param_name }}') in the
|
||||
docstring of the decorated function.
|
||||
|
||||
Args:
|
||||
param_docs: A `ParamDocs` instance or `dict`.
|
||||
|
||||
Returns:
|
||||
A decorator to apply the formatting.
|
||||
|
||||
.. code-block:: text
|
||||
:caption: Example
|
||||
|
||||
>>> param_docs = {"p1": "doc1", "p2": "doc2
|
||||
doc2 second line"}
|
||||
>>> @format_docstring(param_docs)
|
||||
... def func(p1, p2):
|
||||
... '''
|
||||
... Args:
|
||||
... p1: {{ p1 }}
|
||||
... p2: {{ p2 }}
|
||||
... '''
|
||||
>>> import textwrap
|
||||
>>> print(textwrap.dedent(func.__doc__).strip())
|
||||
|
||||
Args:
|
||||
p1: doc1
|
||||
p2: doc2
|
||||
doc2 second line
|
||||
"""
|
||||
param_docs = ParamDocs(param_docs)
|
||||
|
||||
def decorator(func):
|
||||
func.__doc__ = param_docs.format_docstring(func.__doc__)
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
# `{{ ... }}` represents a placeholder.
|
||||
LOG_MODEL_PARAM_DOCS = ParamDocs({
|
||||
"name": "Model name.",
|
||||
"conda_env": (
|
||||
"""Either a dictionary representation of a Conda environment or the path to a conda
|
||||
environment yaml file. If provided, this describes the environment this model should be run in.
|
||||
At a minimum, it should specify the dependencies contained in `get_default_conda_env()`.
|
||||
If ``None``, a conda environment with pip requirements inferred by
|
||||
:func:`mlflow.models.infer_pip_requirements` is added
|
||||
to the model. If the requirement inference fails, it falls back to using
|
||||
`get_default_pip_requirements`. pip requirements from ``conda_env`` are written to a pip
|
||||
``requirements.txt`` file and the full conda environment is written to ``conda.yaml``.
|
||||
The following is an *example* dictionary representation of a conda environment::
|
||||
|
||||
{
|
||||
"name": "mlflow-env",
|
||||
"channels": ["conda-forge"],
|
||||
"dependencies": [
|
||||
"python=3.8.15",
|
||||
{
|
||||
"pip": [
|
||||
"{{ package_name }}==x.y.z"
|
||||
],
|
||||
},
|
||||
],
|
||||
}"""
|
||||
),
|
||||
"pip_requirements": (
|
||||
"""Either an iterable of pip requirement strings
|
||||
(e.g. ``["{{ package_name }}", "-r requirements.txt", "-c constraints.txt"]``) or the string path to
|
||||
a pip requirements file on the local filesystem (e.g. ``"requirements.txt"``). If provided, this
|
||||
describes the environment this model should be run in. If ``None``, a default list of requirements
|
||||
is inferred by :func:`mlflow.models.infer_pip_requirements` from the current software environment.
|
||||
If the requirement inference fails, it falls back to using `get_default_pip_requirements`.
|
||||
Both requirements and constraints are automatically parsed and written to ``requirements.txt`` and
|
||||
``constraints.txt`` files, respectively, and stored as part of the model. Requirements are also
|
||||
written to the ``pip`` section of the model's conda environment (``conda.yaml``) file."""
|
||||
),
|
||||
"extra_pip_requirements": (
|
||||
"""Either an iterable of pip
|
||||
requirement strings
|
||||
(e.g. ``["pandas", "-r requirements.txt", "-c constraints.txt"]``) or the string path to
|
||||
a pip requirements file on the local filesystem (e.g. ``"requirements.txt"``). If provided, this
|
||||
describes additional pip requirements that are appended to a default set of pip requirements
|
||||
generated automatically based on the user's current software environment. Both requirements and
|
||||
constraints are automatically parsed and written to ``requirements.txt`` and ``constraints.txt``
|
||||
files, respectively, and stored as part of the model. Requirements are also written to the ``pip``
|
||||
section of the model's conda environment (``conda.yaml``) file.
|
||||
|
||||
.. warning::
|
||||
The following arguments can't be specified at the same time:
|
||||
|
||||
- ``conda_env``
|
||||
- ``pip_requirements``
|
||||
- ``extra_pip_requirements``
|
||||
|
||||
`This example <https://github.com/mlflow/mlflow/blob/master/examples/pip_requirements/pip_requirements.py>`_ demonstrates how to specify pip requirements using
|
||||
``pip_requirements`` and ``extra_pip_requirements``.""" # noqa: E501
|
||||
),
|
||||
"signature": (
|
||||
"""an instance of the :py:class:`ModelSignature <mlflow.models.ModelSignature>`
|
||||
class that describes the model's inputs and outputs. If not specified but an
|
||||
``input_example`` is supplied, a signature will be automatically inferred
|
||||
based on the supplied input example and model. To disable automatic signature
|
||||
inference when providing an input example, set ``signature`` to ``False``.
|
||||
To manually infer a model signature, call
|
||||
:py:func:`infer_signature() <mlflow.models.infer_signature>` on datasets
|
||||
with valid model inputs, such as a training dataset with the target column
|
||||
omitted, and valid model outputs, like model predictions made on the training
|
||||
dataset, for example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mlflow.models import infer_signature
|
||||
|
||||
train = df.drop_column("target_label")
|
||||
predictions = ... # compute model predictions
|
||||
signature = infer_signature(train, predictions)
|
||||
"""
|
||||
),
|
||||
"metadata": ("Custom metadata dictionary passed to the model and stored in the MLmodel file."),
|
||||
"input_example": (
|
||||
"""one or several instances of valid model input. The input example is used
|
||||
as a hint of what data to feed the model. It will be converted to a Pandas
|
||||
DataFrame and then serialized to json using the Pandas split-oriented
|
||||
format, or a numpy array where the example will be serialized to json
|
||||
by converting it to a list. Bytes are base64-encoded. When the ``signature`` parameter is
|
||||
``None``, the input example is used to infer a model signature.
|
||||
"""
|
||||
),
|
||||
"prompt_template": (
|
||||
"""A string that, if provided, will be used to format the user's input prior
|
||||
to inference. The string should contain a single placeholder, ``{prompt}``, which will be
|
||||
replaced with the user's input. For example: ``"Answer the following question. Q: {prompt} A:"``.
|
||||
|
||||
Currently, only the following pipeline types are supported:
|
||||
|
||||
- `feature-extraction <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.FeatureExtractionPipeline>`_
|
||||
- `fill-mask <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.FillMaskPipeline>`_
|
||||
- `summarization <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.SummarizationPipeline>`_
|
||||
- `text2text-generation <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.Text2TextGenerationPipeline>`_
|
||||
- `text-generation <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.TextGenerationPipeline>`_
|
||||
|
||||
The following example shows how to log a text-generation pipeline with a prompt template and
|
||||
use it via the ``python_function`` (pyfunc) flavor:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import mlflow
|
||||
from transformers import pipeline
|
||||
|
||||
# Initialize a text-generation pipeline
|
||||
generator = pipeline("text-generation", model="gpt2")
|
||||
|
||||
# Define a prompt template. The ``{prompt}`` placeholder will be replaced
|
||||
# with the raw user input at inference time.
|
||||
prompt_template = "Answer the following question concisely.\\n\\nQ: {prompt}\\nA:"
|
||||
|
||||
example_prompt = "What is MLflow?"
|
||||
|
||||
# Log the model with the prompt template and an input example
|
||||
with mlflow.start_run():
|
||||
model_info = mlflow.transformers.log_model(
|
||||
transformers_model=generator,
|
||||
name="qa_text_generator",
|
||||
prompt_template=prompt_template,
|
||||
input_example=example_prompt,
|
||||
)
|
||||
|
||||
# Load the model back as a pyfunc model
|
||||
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
||||
|
||||
# The input to ``predict`` is the raw question string; the prompt template
|
||||
# is applied internally before calling the underlying transformers pipeline.
|
||||
loaded_model.predict("What is experiment tracking?")
|
||||
"""
|
||||
),
|
||||
"code_paths": (
|
||||
"""A list of local filesystem paths to Python file dependencies (or directories
|
||||
containing file dependencies). These files are *prepended* to the system path when the model
|
||||
is loaded. Files declared as dependencies for a given model should have relative
|
||||
imports declared from a common root path if multiple files are defined with import dependencies
|
||||
between them to avoid import errors when loading the model.
|
||||
|
||||
For a detailed explanation of ``code_paths`` functionality, recommended usage patterns and
|
||||
limitations, see the
|
||||
`code_paths usage guide <https://mlflow.org/docs/latest/model/dependencies.html?highlight=code_paths#saving-extra-code-with-an-mlflow-model>`_.
|
||||
"""
|
||||
),
|
||||
"extra_files": (
|
||||
"""A list containing the paths to corresponding extra files, if ``None``, no
|
||||
extra files are added to the model. Remote URIs are resolved to absolute filesystem
|
||||
paths. For example, consider the following ``extra_files`` list:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
extra_files = ["s3://my-bucket/path/to/my_file1", "/local-path/to/my_file2"]
|
||||
|
||||
In this case, the ``"my_file1"`` extra file is downloaded from S3.
|
||||
Model paths will be ["extra_files/my_file1", "extra_files/my_file2"] in the model directory.
|
||||
"""
|
||||
),
|
||||
# Only pyfunc flavor supports `infer_code_paths`.
|
||||
"code_paths_pyfunc": (
|
||||
"""A list of local filesystem paths to Python file dependencies (or directories
|
||||
containing file dependencies). These files are *prepended* to the system path when the model
|
||||
is loaded. Files declared as dependencies for a given model should have relative
|
||||
imports declared from a common root path if multiple files are defined with import dependencies
|
||||
between them to avoid import errors when loading the model.
|
||||
|
||||
You can leave ``code_paths`` argument unset but set ``infer_code_paths`` to ``True`` to let MLflow
|
||||
infer the model code paths. See ``infer_code_paths`` argument doc for details.
|
||||
|
||||
For a detailed explanation of ``code_paths`` functionality, recommended usage patterns and
|
||||
limitations, see the
|
||||
`code_paths usage guide <https://mlflow.org/docs/latest/model/dependencies.html?highlight=code_paths#saving-extra-code-with-an-mlflow-model>`_.
|
||||
"""
|
||||
),
|
||||
"infer_code_paths": (
|
||||
"""If set to ``True``, MLflow automatically infers model code paths. The inferred
|
||||
code path files only include necessary python module files. Only python code files
|
||||
under current working directory are automatically inferable. Default value is
|
||||
``False``.
|
||||
|
||||
.. warning::
|
||||
Please ensure that the custom python module code does not contain sensitive data such as
|
||||
credential token strings, otherwise they might be included in the automatic inferred code
|
||||
path files and be logged to MLflow artifact repository.
|
||||
|
||||
If your custom python module depends on non-python files (e.g. a JSON file) with a relative
|
||||
path to the module code file path, the non-python files can't be automatically inferred as the
|
||||
code path file. To address this issue, you should put all used non-python files outside
|
||||
your custom code directory.
|
||||
|
||||
If a python code file is loaded as the python ``__main__`` module, then this code file can't be
|
||||
inferred as the code path file. If your model depends on classes / functions defined in
|
||||
``__main__`` module, you should use `cloudpickle` to dump your model instance in order to pickle
|
||||
classes / functions in ``__main__``.
|
||||
|
||||
.. Note:: Experimental: This parameter may change or be removed in a future release without warning.
|
||||
"""
|
||||
),
|
||||
"save_pretrained": (
|
||||
"""If set to ``False``, MLflow will not save the Transformer model weight files,
|
||||
instead only saving the reference to the HuggingFace Hub model repository and its commit hash.
|
||||
This is useful when you load the pretrained model from HuggingFace Hub and want to log or save
|
||||
it to MLflow without modifying the model weights. In such case, specifying this flag to
|
||||
``False`` will save the storage space and reduce time to save the model. Please refer to the
|
||||
`Storage-Efficient Model Logging
|
||||
<../../llms/transformers/large-models.html#transformers-save-pretrained-guide>`_ for more detailed
|
||||
usage.
|
||||
|
||||
|
||||
.. warning::
|
||||
|
||||
If the model is saved with ``save_pretrained`` set to ``False``, the model cannot be
|
||||
registered to the MLflow Model Registry. In order to convert the model to the one that
|
||||
can be registered, you can use :py:func:`mlflow.transformers.persist_pretrained_model()`
|
||||
to download the model weights from the HuggingFace Hub and save it in the existing model
|
||||
artifacts. Please refer to `Transformers flavor documentation
|
||||
<../../llms/transformers/large-models.html#persist-pretrained-guide>`_
|
||||
for more detailed usage.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import mlflow.transformers
|
||||
|
||||
model_uri = "YOUR_MODEL_URI_LOGGED_WITH_SAVE_PRETRAINED_FALSE"
|
||||
model = mlflow.transformers.persist_pretrained_model(model_uri)
|
||||
mlflow.register_model(model_uri, "model_name")
|
||||
|
||||
.. important::
|
||||
|
||||
When you save the `PEFT <https://huggingface.co/docs/peft/en/index>`_ model, MLflow will
|
||||
override the `save_pretrained` flag to `False` and only store the PEFT adapter weights. The
|
||||
base model weights are not saved but the reference to the HuggingFace repository and
|
||||
its commit hash are logged instead.
|
||||
"""
|
||||
),
|
||||
"auth_policy": (
|
||||
"""Specifies the authentication policy for the model, which includes two key components.
|
||||
Note that only one of `auth_policy` or `resources` should be defined.
|
||||
|
||||
- **System Auth Policy**: A list of resources required to serve this model.
|
||||
- **User Auth Policy**: A minimal list of scopes that the user should have access to
|
||||
,in order to invoke this model.
|
||||
|
||||
.. Note::
|
||||
Experimental: This parameter may change or be removed in a future release without warning.
|
||||
"""
|
||||
),
|
||||
"params": "A dictionary of parameters to log with the model.",
|
||||
"tags": "A dictionary of tags to log with the model.",
|
||||
"model_type": "The type of the model.",
|
||||
"step": "The step at which to log the model outputs and metrics",
|
||||
"model_id": "The ID of the model.",
|
||||
"prompts": """\
|
||||
A list of prompt URIs registered in the MLflow Prompt Registry, to be associated with the model.
|
||||
Each prompt URI should be in the form ``prompt:/<name>/<version>``. The prompts should be
|
||||
registered in the MLflow Prompt Registry before being associated with the model.
|
||||
|
||||
This will create a mutual link between the model and the prompt. The associated prompts can be
|
||||
seen in the model's metadata stored in the MLmodel file. From the Prompt Registry UI, you can
|
||||
navigate to the model as well.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import mlflow
|
||||
|
||||
prompt_template = "Hi, {name}! How are you doing today?"
|
||||
|
||||
# Register a prompt in the MLflow Prompt Registry
|
||||
mlflow.prompts.register_prompt("my_prompt", prompt_template, description="A simple prompt")
|
||||
|
||||
# Log a model with the registered prompt
|
||||
with mlflow.start_run():
|
||||
model_info = mlflow.pyfunc.log_model(
|
||||
name=MyModel(),
|
||||
name="model",
|
||||
prompts=["prompt:/my_prompt/1"]
|
||||
)
|
||||
|
||||
print(model_info.prompts)
|
||||
# Output: ['prompt:/my_prompt/1']
|
||||
|
||||
# Load the prompt
|
||||
prompt = mlflow.genai.load_prompt(model_info.prompts[0])
|
||||
""",
|
||||
})
|
||||
|
||||
|
||||
def get_module_min_and_max_supported_ranges(flavor_name):
|
||||
"""
|
||||
Extracts the minimum and maximum supported package versions from the provided module name.
|
||||
The version information is provided via the yaml-to-python-script generation script in
|
||||
dev/update_ml_package_versions.py which writes a python file to the importable namespace of
|
||||
mlflow.ml_package_versions
|
||||
|
||||
Args:
|
||||
flavor_name: The flavor name registered in ml_package_versions.py
|
||||
|
||||
Returns:
|
||||
tuple of module name, minimum supported version, maximum supported version as strings.
|
||||
"""
|
||||
if flavor_name == "pyspark.ml":
|
||||
# pyspark.ml is a special case of spark flavor
|
||||
flavor_name = "spark"
|
||||
|
||||
module_name = _ML_PACKAGE_VERSIONS[flavor_name]["package_info"].get("module_name", flavor_name)
|
||||
versions = _ML_PACKAGE_VERSIONS[flavor_name]["models"]
|
||||
min_version = versions["minimum"]
|
||||
max_version = versions["maximum"]
|
||||
return module_name, min_version, max_version
|
||||
|
||||
|
||||
def _do_version_compatibility_warning(msg: str):
|
||||
"""
|
||||
Isolate the warn call to show the warning only once.
|
||||
"""
|
||||
warnings.warn(msg, category=UserWarning, stacklevel=2)
|
||||
|
||||
|
||||
def docstring_version_compatibility_warning(integration_name):
|
||||
"""
|
||||
Generates a docstring that can be applied as a note stating a version compatibility range for
|
||||
a given flavor and optionally raises a warning if the installed version is outside of the
|
||||
supported range.
|
||||
|
||||
Args:
|
||||
integration_name: The name of the module as stored within ml-package-versions.yml
|
||||
|
||||
Returns:
|
||||
The wrapped function with the additional docstring header applied
|
||||
"""
|
||||
|
||||
def annotated_func(func):
|
||||
# NB: if using this decorator, ensure the package name to module name reference is
|
||||
# updated with the flavor's `save` and `load` functions being used within
|
||||
# ml-package-version.yml file.
|
||||
min_ver, max_ver, pip_release = get_min_max_version_and_pip_release(
|
||||
integration_name, "models"
|
||||
)
|
||||
notice = (
|
||||
f"The '{integration_name}' MLflow Models integration is known to be compatible with "
|
||||
f"``{min_ver}`` <= ``{pip_release}`` <= ``{max_ver}``. "
|
||||
f"MLflow Models integrations with {integration_name} may not succeed when used with "
|
||||
"package versions outside of this range."
|
||||
)
|
||||
|
||||
func.__doc__ = (
|
||||
" .. Note:: " + notice + "\n" * 2 + func.__doc__ if func.__doc__ else notice
|
||||
)
|
||||
|
||||
return func
|
||||
|
||||
return annotated_func
|
||||
Reference in New Issue
Block a user