python - 将 MultiLabelBinarizer 打包到 scikit-learn Pipeline 中以推理新数据

Question

我正在构建一个多标签分类器来根据文本字段预测标签。例如，根据电影标题预测类型。我想使用 MultiLabelBinarizer() 对包含所有适用流派标签的列进行二值化。例如，['action','comedy','drama'] 被分成具有 0/1 值的三列。

我使用 MultiLabelBinarizer() 的原因是我可以使用内置的 inverse_transform() 函数来转换输出数组(例如 array([0, 0, 1, 0, 1]) 直接转换为用户友好的文本输出 (['action','drama'])。

分类器可以工作，但我在预测新数据时遇到问题。我找不到将 MultiLabelBinarizer() 集成到我的管道中的方法，以便可以保存并重新加载它以推断新数据。一种解决方案是将其单独保存为 pickle 对象并每次加载回来，但我想避免在生产中出现这种依赖关系。

我知道这与我在管道中内置的 tf-idf 向量类似，但不同之处在于它应用于目标列(流派标签)而不是我的自变量(文本注释)。这是我训练多标签 SVM 的代码:

def svm_train(df):  
  mlb = MultiLabelBinarizer()
  y = mlb.fit_transform(df['Genres'])

  with mlflow.start_run():
    x_train, x_test, y_train, y_test = train_test_split(df['Movie Title'], y, test_size=0.3)

    # Instantiate TF-IDF Vectorizer and SVM Model
    tfidf_vect = TfidfVectorizer()
    mdl = OneVsRestClassifier(LinearSVC(loss='hinge'))
    svm_pipeline = Pipeline([('tfidf', tfidf_vect), ('clf', mdl)])

    svm_pipeline.fit(x_train, y_train)
    prediction = svm_pipeline.predict(x_test)

    report = classification_report(y_test, prediction, target_names=mlb.classes_)

    mlflow.sklearn.log_model(svm_pipeline, "Multilabel Classifier")
    mlflow.log_artifact(mlb, "MLB")

  return(report)

svm_train(df)

推理包括在单独的 Databricks 笔记本中从 MLflow 重新加载保存的模型(与加载回 pickle 文件相同)并使用管道进行预测:

def predict_labels(new_data):
  model_uri = '...MLflow path...'
  model = mlflow.sklearn.load_model(model_uri)
  predictions = model.predict(new_data)
  # If I can't package the MultiLabelBinarizer() into the Pipeline, this 
  # is where I'd have to load the pickle object mlb
  # so that I can inverse_transform()
  return mlb.inverse_transform(predictions)

new_data = ['Some movie title']
predict_labels(new_data)

['action','comedy']

这是我正在使用的所有库:

import pandas as pd
import numpy as np
import mlflow
import mlflow.sklearn
import glob, os
from pyspark.sql import DataFrame
from sklearn.pipeline import Pipeline
from sklearn import preprocessing
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.multiclass import OneVsRestClassifier
from sklearn import svm
from sklearn.svm import LinearSVC
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from sklearn.metrics import accuracy_score, precision_score, recall_score

Answer 1

对于您的用例，您可能需要考虑使用 MLflow's functionality for persisting custom models 。根据the docs :

While MLflow’s built-in model persistence utilities are convenient for packaging models from various popular ML libraries in MLflow Model format, they do not cover every use case. For example, you may want to use a model from an ML library that is not explicitly supported by MLflow’s built-in flavors. Alternatively, you may want to package custom inference code and data to create an MLflow Model. Fortunately, MLflow provides two solutions that can be used to accomplish these tasks: Custom Python Models and Custom Flavors.

特别是，您应该能够以类似于链接示例中的 XGBoost 模型的方式将 MultiLabelIndexer 作为工件与 Sklearn 模型一起记录，然后在预测时将其加载回来，如下所示:

# Save sklearn model & multilabel indexer to paths on the local filesystem
sklearn_model_path = "some/local/path"
labelindexer_path = "another/local/path"
# ... save your models objects here to sklearn_model_path and labelindexer_path

# Define the custom model class
import mlflow.pyfunc
class SklearnWrapper(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        import pickle, mlflow
        with open(context["indexer_path"], 'rb') as handle:
            self.indexer = pickle.load(handle)
        self.pipeline = mlflow.sklearn.load_model("pipeline_path")

    def predict(self, context, model_input):
        pipeline_preds = self.pipeline.predict(model_input)
        return self.indexer.inverse_transform(pipeline_preds)

# Create a Conda environment for the new MLflow Model that contains the XGBoost library
# as a dependency, as well as the required CloudPickle library
import cloudpickle
import sklearn
conda_env = {
    'channels': ['defaults'],
    'dependencies': [
      'sklearn={}'.format(sklearn.__version__),
      'cloudpickle={}'.format(cloudpickle.__version__),
    ],
    'name': 'sklearn_env'
}

# Save the MLflow Model
artifacts = {
    "pipeline_path": sklearn_model_path,
    "indexer_path": labelindexer_path,
}
mlflow_pyfunc_model_path = "sklearn_mlflow_pyfunc"
mlflow.pyfunc.save_model(
        path=mlflow_pyfunc_model_path, python_model=XGBWrapper(), artifacts=artifacts,
        conda_env=conda_env)

# Load the model in `python_function` format
loaded_model = mlflow.pyfunc.load_model(mlflow_pyfunc_model_path)
# Predict on a pandas DataFrame
import pandas as pd
loaded_model.predict(pd.DataFrame(...))

请注意，我们的自定义模型仍会加载回 MultiLabelIndexer，但 MLflow 会将索引器与您的管道和自定义模型逻辑一起保留，以便您可以将模型视为用于生产部署的单个连贯单元。