528 lines
19 KiB
Python
528 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import warnings
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Union
|
|
|
|
import cloudpickle
|
|
import yaml
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.langchain.utils import (
|
|
_BASE_LOAD_KEY,
|
|
_CONFIG_LOAD_KEY,
|
|
_MODEL_DATA_FOLDER_NAME,
|
|
_MODEL_DATA_KEY,
|
|
_MODEL_DATA_PKL_FILE_NAME,
|
|
_MODEL_DATA_YAML_FILE_NAME,
|
|
_MODEL_LOAD_KEY,
|
|
_MODEL_TYPE_KEY,
|
|
_PICKLE_LOAD_KEY,
|
|
_RUNNABLE_LOAD_KEY,
|
|
_load_base_lcs,
|
|
_load_from_json,
|
|
_load_from_pickle,
|
|
_load_from_yaml,
|
|
_patch_loader,
|
|
_save_base_lcs,
|
|
_validate_and_prepare_lc_model_or_path,
|
|
base_lc_types,
|
|
custom_type_to_loader_dict,
|
|
get_unsupported_model_message,
|
|
lc_runnable_assign_types,
|
|
lc_runnable_binding_types,
|
|
lc_runnable_branch_types,
|
|
lc_runnable_with_steps_types,
|
|
lc_runnables_types,
|
|
patch_langchain_type_to_cls_dict,
|
|
picklable_runnable_types,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from langchain.schema.runnable import Runnable
|
|
|
|
_STEPS_FOLDER_NAME = "steps"
|
|
_RUNNABLE_STEPS_FILE_NAME = "steps.yaml"
|
|
_BRANCHES_FOLDER_NAME = "branches"
|
|
_MAPPER_FOLDER_NAME = "mapper"
|
|
_RUNNABLE_BRANCHES_FILE_NAME = "branches.yaml"
|
|
_DEFAULT_BRANCH_NAME = "default"
|
|
_RUNNABLE_BINDING_CONF_FILE_NAME = "binding_conf.yaml"
|
|
|
|
|
|
@patch_langchain_type_to_cls_dict
|
|
def _load_model_from_config(path, model_config):
|
|
from langchain.chains.loading import type_to_loader_dict as chains_type_to_loader_dict
|
|
from langchain.llms import get_type_to_cls_dict as llms_get_type_to_cls_dict
|
|
|
|
try:
|
|
from langchain.prompts.loading import type_to_loader_dict as prompts_types
|
|
except ImportError:
|
|
prompts_types = {"prompt", "few_shot_prompt"}
|
|
|
|
config_path = os.path.join(path, model_config.get(_MODEL_DATA_KEY, _MODEL_DATA_YAML_FILE_NAME))
|
|
# Load runnables from config file
|
|
if config_path.endswith(".yaml"):
|
|
config = _load_from_yaml(config_path)
|
|
elif config_path.endswith(".json"):
|
|
config = _load_from_json(config_path)
|
|
else:
|
|
raise MlflowException(
|
|
f"Cannot load runnable without a config file. Got path {config_path}."
|
|
)
|
|
_type = config.get("_type")
|
|
if _type in chains_type_to_loader_dict:
|
|
from langchain.chains.loading import load_chain
|
|
|
|
return _patch_loader(load_chain)(config_path)
|
|
elif _type in prompts_types:
|
|
from langchain.prompts.loading import load_prompt
|
|
|
|
return load_prompt(config_path)
|
|
elif _type in llms_get_type_to_cls_dict():
|
|
from langchain_community.llms.loading import load_llm
|
|
|
|
return _patch_loader(load_llm)(config_path)
|
|
elif _type in custom_type_to_loader_dict():
|
|
return custom_type_to_loader_dict()[_type](config)
|
|
raise MlflowException(f"Unsupported type {_type} for loading.")
|
|
|
|
|
|
def _load_model_from_path(path: str, model_config=None):
|
|
model_load_fn = model_config.get(_MODEL_LOAD_KEY)
|
|
if model_load_fn == _RUNNABLE_LOAD_KEY:
|
|
return _load_runnables(path, model_config)
|
|
if model_load_fn == _BASE_LOAD_KEY:
|
|
return _load_base_lcs(path, model_config)
|
|
if model_load_fn == _CONFIG_LOAD_KEY:
|
|
return _load_model_from_config(path, model_config)
|
|
if model_load_fn == _PICKLE_LOAD_KEY:
|
|
return _load_from_pickle(os.path.join(path, model_config.get(_MODEL_DATA_KEY)))
|
|
raise MlflowException(f"Unsupported model load key {model_load_fn}")
|
|
|
|
|
|
def _validate_path(file_path: Union[str, Path]):
|
|
load_path = Path(file_path)
|
|
if not load_path.exists() or not load_path.is_dir():
|
|
raise MlflowException(
|
|
f"Path {load_path} must be an existing directory in order to load model."
|
|
)
|
|
return load_path
|
|
|
|
|
|
def _load_runnable_with_steps(file_path: Union[Path, str], model_type: str):
|
|
"""Load the model
|
|
|
|
Args:
|
|
file_path: Path to file to load the model from.
|
|
model_type: Type of the model to load.
|
|
"""
|
|
from langchain.schema.runnable import RunnableParallel, RunnableSequence
|
|
|
|
load_path = _validate_path(file_path)
|
|
|
|
steps_conf_file = load_path / _RUNNABLE_STEPS_FILE_NAME
|
|
if not steps_conf_file.exists():
|
|
raise MlflowException(
|
|
f"File {steps_conf_file} must exist in order to load runnable with steps."
|
|
)
|
|
steps_conf = _load_from_yaml(steps_conf_file)
|
|
steps_path = load_path / _STEPS_FOLDER_NAME
|
|
_validate_path(steps_path)
|
|
|
|
steps = {}
|
|
# ignore hidden files
|
|
for step in (f for f in os.listdir(steps_path) if not f.startswith(".")):
|
|
config = steps_conf.get(step)
|
|
# load model from the folder of the step
|
|
runnable = _load_model_from_path(os.path.join(steps_path, step), config)
|
|
steps[step] = runnable
|
|
|
|
if model_type == RunnableSequence.__name__:
|
|
steps = [value for _, value in sorted(steps.items(), key=lambda item: int(item[0]))]
|
|
return runnable_sequence_from_steps(steps)
|
|
if model_type == RunnableParallel.__name__:
|
|
return RunnableParallel(steps)
|
|
|
|
|
|
def runnable_sequence_from_steps(steps):
|
|
"""Construct a RunnableSequence from steps.
|
|
|
|
Args:
|
|
steps: List of steps to construct the RunnableSequence from.
|
|
"""
|
|
from langchain.schema.runnable import RunnableSequence
|
|
|
|
if len(steps) < 2:
|
|
raise ValueError(f"RunnableSequence must have at least 2 steps, got {len(steps)}.")
|
|
|
|
first, *middle, last = steps
|
|
return RunnableSequence(first=first, middle=middle, last=last)
|
|
|
|
|
|
def _load_runnable_branch(file_path: Union[Path, str]):
|
|
"""Load the model
|
|
|
|
Args:
|
|
file_path: Path to file to load the model from.
|
|
"""
|
|
from langchain.schema.runnable import RunnableBranch
|
|
|
|
load_path = _validate_path(file_path)
|
|
|
|
branches_conf_file = load_path / _RUNNABLE_BRANCHES_FILE_NAME
|
|
if not branches_conf_file.exists():
|
|
raise MlflowException(
|
|
f"File {branches_conf_file} must exist in order to load runnable with steps."
|
|
)
|
|
branches_conf = _load_from_yaml(branches_conf_file)
|
|
branches_path = load_path / _BRANCHES_FOLDER_NAME
|
|
_validate_path(branches_path)
|
|
|
|
branches = []
|
|
for branch in os.listdir(branches_path):
|
|
# load model from the folder of the branch
|
|
if branch == _DEFAULT_BRANCH_NAME:
|
|
default_branch_path = branches_path / _DEFAULT_BRANCH_NAME
|
|
default = _load_model_from_path(
|
|
default_branch_path, branches_conf.get(_DEFAULT_BRANCH_NAME)
|
|
)
|
|
else:
|
|
branch_tuple = []
|
|
for i in range(2):
|
|
config = branches_conf.get(f"{branch}-{i}")
|
|
runnable = _load_model_from_path(
|
|
os.path.join(branches_path, branch, str(i)), config
|
|
)
|
|
branch_tuple.append(runnable)
|
|
branches.append(tuple(branch_tuple))
|
|
|
|
# default branch must be the last branch
|
|
branches.append(default)
|
|
|
|
return RunnableBranch(*branches)
|
|
|
|
|
|
def _load_runnable_assign(file_path: Union[Path, str]):
|
|
"""Load the model
|
|
|
|
Args:
|
|
file_path: Path to file to load the model from.
|
|
"""
|
|
from langchain.schema.runnable.passthrough import RunnableAssign
|
|
|
|
load_path = _validate_path(file_path)
|
|
|
|
mapper_file = load_path / _MAPPER_FOLDER_NAME
|
|
_validate_path(mapper_file)
|
|
mapper = _load_runnable_with_steps(mapper_file, "RunnableParallel")
|
|
return RunnableAssign(mapper)
|
|
|
|
|
|
def _load_runnable_binding(file_path: Union[Path, str]):
|
|
"""
|
|
Load runnable binding model from the path
|
|
"""
|
|
from langchain.schema.runnable import RunnableBinding
|
|
|
|
load_path = _validate_path(file_path)
|
|
|
|
model_conf = _load_from_yaml(load_path / _RUNNABLE_BINDING_CONF_FILE_NAME)
|
|
for field, value in model_conf.items():
|
|
if _is_json_primitive(value):
|
|
model_conf[field] = value
|
|
# value is dictionary
|
|
else:
|
|
model_conf[field] = _load_model_from_path(load_path, value)
|
|
return RunnableBinding(**model_conf)
|
|
|
|
|
|
def _save_internal_runnables(runnable, path, loader_fn, persist_dir):
|
|
conf = {}
|
|
if isinstance(runnable, lc_runnables_types()):
|
|
conf[_MODEL_TYPE_KEY] = runnable.__class__.__name__
|
|
conf.update(_save_runnables(runnable, path, loader_fn, persist_dir))
|
|
elif isinstance(runnable, base_lc_types()):
|
|
lc_model = _validate_and_prepare_lc_model_or_path(runnable, loader_fn)
|
|
conf[_MODEL_TYPE_KEY] = lc_model.__class__.__name__
|
|
conf.update(_save_base_lcs(lc_model, path, loader_fn, persist_dir))
|
|
else:
|
|
conf = {
|
|
_MODEL_TYPE_KEY: runnable.__class__.__name__,
|
|
_MODEL_DATA_KEY: _MODEL_DATA_YAML_FILE_NAME,
|
|
_MODEL_LOAD_KEY: _CONFIG_LOAD_KEY,
|
|
}
|
|
model_path = path / _MODEL_DATA_YAML_FILE_NAME
|
|
|
|
_warning_if_imported_from_lc_partner_pkg(runnable)
|
|
|
|
# Save some simple runnables that langchain natively supports.
|
|
if hasattr(runnable, "save"):
|
|
runnable.save(model_path)
|
|
elif hasattr(runnable, "dict"):
|
|
runnable_dict = runnable.dict()
|
|
with open(model_path, "w") as f:
|
|
yaml.dump(runnable_dict, f, default_flow_style=False)
|
|
# if the model cannot be loaded back, then `dict` is not enough for saving.
|
|
_load_model_from_config(path, conf)
|
|
else:
|
|
raise Exception("Cannot save runnable without `save` or `dict` methods.")
|
|
return conf
|
|
|
|
|
|
_LC_PARTNER_MODULE_PATTERN = re.compile(
|
|
r"langchain_(?!core|community|experimental|cli|text-splitters)([a-z0-9-]+)$"
|
|
)
|
|
|
|
|
|
def _warning_if_imported_from_lc_partner_pkg(runnable):
|
|
"""
|
|
Issues a warning if the model contains LangChain partner packages in its requirements.
|
|
|
|
Popular integrations like OpenAI have been migrated from the central langchain-community
|
|
package to their own partner packages (e.g. langchain-openai). However, the class loading
|
|
mechanism in MLflow does not handle partner packages and always loads the community version.
|
|
This can lead to unexpected behavior because the community version is no longer maintained.
|
|
"""
|
|
module = runnable.__module__
|
|
root_module = module.split(".")[0]
|
|
if m := _LC_PARTNER_MODULE_PATTERN.match(root_module):
|
|
warnings.warn(
|
|
"Your model contains a class imported from the LangChain partner package "
|
|
f"`langchain-{m.group(1)}`. When loading the model back, MLflow will use the "
|
|
"community version of the classes instead of the partner packages, which may "
|
|
"lead to unexpected behavior. To ensure that the model is loaded correctly, "
|
|
"it is recommended to save the model with the 'model-from-code' method "
|
|
"instead: https://mlflow.org/docs/latest/models.html#models-from-code"
|
|
)
|
|
|
|
|
|
def _save_runnable_with_steps(model, file_path: Union[Path, str], loader_fn=None, persist_dir=None): # noqa: D417
|
|
"""Save the model with steps. Currently it supports saving RunnableSequence and
|
|
RunnableParallel.
|
|
|
|
If saving a RunnableSequence, steps is a list of Runnable objects. We save each step to the
|
|
subfolder named by the step index.
|
|
e.g. - model
|
|
- steps
|
|
- 0
|
|
- model.yaml
|
|
- 1
|
|
- model.pkl
|
|
- steps.yaml
|
|
If saving a RunnableParallel, steps is a dictionary of key-Runnable pairs. We save each step to
|
|
the subfolder named by the key.
|
|
e.g. - model
|
|
- steps
|
|
- context
|
|
- model.yaml
|
|
- question
|
|
- model.pkl
|
|
- steps.yaml
|
|
|
|
We save steps.yaml file to the model folder. It contains each step's model's configuration.
|
|
|
|
Args:
|
|
model: Runnable to be saved.
|
|
file_path: Path to file to save the model to.
|
|
"""
|
|
# Convert file to Path object.
|
|
save_path = Path(file_path)
|
|
save_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Save steps into a folder
|
|
steps_path = save_path / _STEPS_FOLDER_NAME
|
|
steps_path.mkdir()
|
|
|
|
steps = get_runnable_steps(model)
|
|
if isinstance(steps, list):
|
|
generator = enumerate(steps)
|
|
elif isinstance(steps, dict):
|
|
generator = steps.items()
|
|
else:
|
|
raise MlflowException(
|
|
f"Runnable {model} steps attribute must be either a list or a dictionary. "
|
|
f"Got {type(steps).__name__}."
|
|
)
|
|
unsaved_runnables = {}
|
|
steps_conf = {}
|
|
for key, runnable in generator:
|
|
step = str(key)
|
|
# Save each step into a subfolder named by step
|
|
save_runnable_path = steps_path / step
|
|
save_runnable_path.mkdir()
|
|
try:
|
|
steps_conf[step] = _save_internal_runnables(
|
|
runnable, save_runnable_path, loader_fn, persist_dir
|
|
)
|
|
except Exception as e:
|
|
unsaved_runnables[step] = f"{runnable.get_name()} -- {e}"
|
|
|
|
if unsaved_runnables:
|
|
raise MlflowException(f"Failed to save runnable sequence: {unsaved_runnables}.")
|
|
|
|
# save steps configs
|
|
with save_path.joinpath(_RUNNABLE_STEPS_FILE_NAME).open("w") as f:
|
|
yaml.dump(steps_conf, f, default_flow_style=False)
|
|
|
|
|
|
def _save_runnable_branch(model, file_path, loader_fn, persist_dir):
|
|
"""
|
|
Save runnable branch in to path.
|
|
"""
|
|
save_path = Path(file_path)
|
|
save_path.mkdir(parents=True, exist_ok=True)
|
|
# save branches into a folder
|
|
branches_path = save_path / _BRANCHES_FOLDER_NAME
|
|
branches_path.mkdir()
|
|
|
|
unsaved_runnables = {}
|
|
branches_conf = {}
|
|
for index, branch_tuple in enumerate(model.branches):
|
|
# Save each branch into a subfolder named by index
|
|
# and save condition and runnable into subfolder
|
|
for i, runnable in enumerate(branch_tuple):
|
|
save_runnable_path = branches_path / str(index) / str(i)
|
|
save_runnable_path.mkdir(parents=True)
|
|
branches_conf[f"{index}-{i}"] = {}
|
|
|
|
try:
|
|
branches_conf[f"{index}-{i}"] = _save_internal_runnables(
|
|
runnable, save_runnable_path, loader_fn, persist_dir
|
|
)
|
|
except Exception as e:
|
|
unsaved_runnables[f"{index}-{i}"] = f"{runnable.get_name()} -- {e}"
|
|
|
|
# save default branch
|
|
default_branch_path = branches_path / _DEFAULT_BRANCH_NAME
|
|
default_branch_path.mkdir()
|
|
try:
|
|
branches_conf[_DEFAULT_BRANCH_NAME] = _save_internal_runnables(
|
|
model.default, default_branch_path, loader_fn, persist_dir
|
|
)
|
|
except Exception as e:
|
|
unsaved_runnables[_DEFAULT_BRANCH_NAME] = f"{model.default.get_name()} -- {e}"
|
|
if unsaved_runnables:
|
|
raise MlflowException(f"Failed to save runnable branch: {unsaved_runnables}.")
|
|
|
|
# save branches configs
|
|
with save_path.joinpath(_RUNNABLE_BRANCHES_FILE_NAME).open("w") as f:
|
|
yaml.dump(branches_conf, f, default_flow_style=False)
|
|
|
|
|
|
def _save_runnable_assign(model, file_path, loader_fn=None, persist_dir=None):
|
|
from langchain.schema.runnable import RunnableParallel
|
|
|
|
save_path = Path(file_path)
|
|
save_path.mkdir(parents=True, exist_ok=True)
|
|
# save mapper into a folder
|
|
mapper_path = save_path / _MAPPER_FOLDER_NAME
|
|
mapper_path.mkdir()
|
|
|
|
if not isinstance(model.mapper, RunnableParallel):
|
|
raise MlflowException(
|
|
f"Failed to save model {model} with type {model.__class__.__name__}. "
|
|
"RunnableAssign's mapper must be a RunnableParallel."
|
|
)
|
|
_save_runnable_with_steps(model.mapper, mapper_path, loader_fn, persist_dir)
|
|
|
|
|
|
def _is_json_primitive(value):
|
|
return (
|
|
value is None
|
|
or isinstance(value, (str, int, float, bool))
|
|
or (isinstance(value, list) and all(_is_json_primitive(v) for v in value))
|
|
)
|
|
|
|
|
|
def _save_runnable_binding(model, file_path, loader_fn=None, persist_dir=None):
|
|
save_path = Path(file_path)
|
|
save_path.mkdir(parents=True, exist_ok=True)
|
|
model_config = {}
|
|
|
|
# runnableBinding bound is the real runnable to be invoked
|
|
model_config["bound"] = _save_internal_runnables(model.bound, save_path, loader_fn, persist_dir)
|
|
|
|
# save other fields
|
|
for field, value in model.dict().items():
|
|
if _is_json_primitive(value):
|
|
model_config[field] = value
|
|
elif field != "bound":
|
|
model_config[field] = {
|
|
_MODEL_LOAD_KEY: _PICKLE_LOAD_KEY,
|
|
_MODEL_DATA_KEY: f"{field}.pkl",
|
|
}
|
|
_pickle_object(value, os.path.join(save_path, f"{field}.pkl"))
|
|
|
|
# save fields configs
|
|
with save_path.joinpath(_RUNNABLE_BINDING_CONF_FILE_NAME).open("w") as f:
|
|
yaml.dump(model_config, f, default_flow_style=False)
|
|
|
|
|
|
def _pickle_object(model, path: str):
|
|
if not path.endswith(".pkl"):
|
|
raise ValueError(f"File path must end with .pkl, got {path}.")
|
|
with open(path, "wb") as f:
|
|
cloudpickle.dump(model, f)
|
|
|
|
|
|
def _save_runnables(model, path, loader_fn=None, persist_dir=None):
|
|
model_data_kwargs = {
|
|
_MODEL_LOAD_KEY: _RUNNABLE_LOAD_KEY,
|
|
_MODEL_TYPE_KEY: model.__class__.__name__,
|
|
}
|
|
if isinstance(model, lc_runnable_with_steps_types()):
|
|
model_data_path = _MODEL_DATA_FOLDER_NAME
|
|
_save_runnable_with_steps(
|
|
model, os.path.join(path, model_data_path), loader_fn, persist_dir
|
|
)
|
|
elif isinstance(model, picklable_runnable_types()):
|
|
model_data_path = _MODEL_DATA_PKL_FILE_NAME
|
|
_pickle_object(model, os.path.join(path, model_data_path))
|
|
elif isinstance(model, lc_runnable_branch_types()):
|
|
model_data_path = _MODEL_DATA_FOLDER_NAME
|
|
_save_runnable_branch(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
|
|
elif isinstance(model, lc_runnable_assign_types()):
|
|
model_data_path = _MODEL_DATA_FOLDER_NAME
|
|
_save_runnable_assign(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
|
|
elif isinstance(model, lc_runnable_binding_types()):
|
|
model_data_path = _MODEL_DATA_FOLDER_NAME
|
|
_save_runnable_binding(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
|
|
else:
|
|
raise MlflowException.invalid_parameter_value(
|
|
get_unsupported_model_message(type(model).__name__)
|
|
)
|
|
model_data_kwargs[_MODEL_DATA_KEY] = model_data_path
|
|
return model_data_kwargs
|
|
|
|
|
|
def _load_runnables(path, conf):
|
|
model_type = conf.get(_MODEL_TYPE_KEY)
|
|
model_data = conf.get(_MODEL_DATA_KEY, _MODEL_DATA_YAML_FILE_NAME)
|
|
if model_type in (x.__name__ for x in lc_runnable_with_steps_types()):
|
|
return _load_runnable_with_steps(os.path.join(path, model_data), model_type)
|
|
if (
|
|
model_type in (x.__name__ for x in picklable_runnable_types())
|
|
or model_data == _MODEL_DATA_PKL_FILE_NAME
|
|
):
|
|
return _load_from_pickle(os.path.join(path, model_data))
|
|
if model_type in (x.__name__ for x in lc_runnable_branch_types()):
|
|
return _load_runnable_branch(os.path.join(path, model_data))
|
|
if model_type in (x.__name__ for x in lc_runnable_assign_types()):
|
|
return _load_runnable_assign(os.path.join(path, model_data))
|
|
if model_type in (x.__name__ for x in lc_runnable_binding_types()):
|
|
return _load_runnable_binding(os.path.join(path, model_data))
|
|
raise MlflowException.invalid_parameter_value(get_unsupported_model_message(model_type))
|
|
|
|
|
|
def get_runnable_steps(model: Runnable):
|
|
try:
|
|
return model.steps
|
|
except AttributeError:
|
|
# RunnableParallel stores steps as `steps__` attribute since version 0.16.0, while it was
|
|
# stored as `steps` attribute before that and other runnables like RunnableSequence still
|
|
# has `steps` property.
|
|
return model.steps__
|