Palantir Foundry模型适配器推理仅返回22条结果技术问询

问题背景 #

正在使用Foundry Model Adapters学习Palantir Foundry建模，作为机器学习新手，参考多标签分类教程搭建模型。当前Foundry栈中模型适配器处于Beta阶段，优先保证模型在Foundry中运行而非精度。

已在Modeling Objectives发布staging模型，但运行模型推理转换时，无论输入何种数据集，仅返回22条结果，预期应返回3000+条。相同代码在Code Workspace测试运行正常，怀疑问题出在适配器配置中。

使用数据集 #

PubMed多标签文本分类数据集

模型训练代码 #

from transforms.api import transform, Input, Output
from palantir_models.transforms import ModelOutput
from palantir_models.models import ModelVersionChangeType
from main.model_adapters.adapter import ExampleModelAdapter
import pandas as pd
import numpy as np
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.multioutput import MultiOutputClassifier
from sklearn.linear_model import LogisticRegression


@transform(
     features_and_labels_input=Input("/Modeling Tutorials/PubMed Multi Label Text Classification Dataset"),
     model_output=ModelOutput("/Modeling Tutorials/multi label classifier/multi_label_classifier"),
 )
 def compute(features_and_labels_input, model_output):
     df = features_and_labels_input.pandas()

     df = df.drop(['Title', 'meshMajor', 'pmid', 'meshid', 'meshroot', 'Z', 'V', 'N'], axis=1)
     X = df["abstractText"]
     y = np.asarray(df[df.columns[1:]])

     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
     clf = MultiOutputClassifier(LogisticRegression())
     pipe = Pipeline([
             ('tfidf', TfidfVectorizer()),
             ('clf', clf)
     ])

     model = pipe.fit(X_train, y_train)

     # Wrap the trained model in a ModelAdapter
     foundry_model = ExampleModelAdapter(model)

     # Publish and write the trained model to Foundry
     model_output.publish(
         model_adapter=foundry_model,
         change_type=ModelVersionChangeType.MINOR   # How to increment model version number: MAJOR (x.0.0), MINOR(0.x.0), PATCH (0.0.x)
     )

Sklearn适配器代码 #

from palantir_models.models import ModelAdapter, PythonEnvironment, CondaDependency
from palantir_models.models.api import ModelApi, ModelApiColumn, ModelInput, ModelOutput, DFType
from palantir_models.models._types import CondaVersionExact
import pickle
import os
import pandas as pd


class ExampleModelAdapter(ModelAdapter):
    MODEL_SAVE_LOCATION = 'model.pkl'
    METADATA_SAVE_LOCATION = 'metadata.pkl'
    PREDICTION_COLUMN_NAME_KEY = 'prediction_column_name'

    model = None
    prediction_column_name = None

    def __init__(self, model, prediction_column_name='prediction'):
        self.model = model
        self.prediction_column_name = prediction_column_name

    @classmethod
    def load(cls, state_reader, container_context):
        with state_reader.extract_to_temp_dir() as tmp_dir:
            model = pickle.load(open(os.path.join(tmp_dir, ExampleModelAdapter.MODEL_SAVE_LOCATION), "rb"))
            metadata = pickle.load(open(os.path.join(tmp_dir, ExampleModelAdapter.METADATA_SAVE_LOCATION), "rb"))
            prediction_column_name = metadata[ExampleModelAdapter.PREDICTION_COLUMN_NAME_KEY]
        return cls(model, prediction_column_name)

    def save(self, state_writer):
        with state_writer.open(ExampleModelAdapter.MODEL_SAVE_LOCATION, "wb") as model_file:
        pickle.dump(self.model, model_file)
        with state_writer.open(ExampleModelAdapter.METADATA_SAVE_LOCATION, "wb") as metadata_file:
            metadata = {
                ExampleModelAdapter.PREDICTION_COLUMN_NAME_KEY: self.prediction_column_name
            }
            pickle.dump(metadata, metadata_file)

    @classmethod
    def api(cls):
        inputs = [
            ModelInput.Tabular(name="input_df",
                           df_type=DFType.PANDAS,
                           columns=[ModelApiColumn(name="features", type=tuple)])
        ]
        outputs = [
            ModelOutput.Tabular(name="output_df",
                            columns=[ModelApiColumn(name="prediction", type=list)])
        ]
        return ModelApi(inputs, outputs)

    def run_inference(self, inputs, outputs):
        df_in = inputs.input_df
        df_out = outputs.output_df
        df = pd.DataFrame(self.model.predict(df_in))
        df_out.write(df)

    @classmethod
    def dependencies(cls):
        # DO NOT MODIFY THIS FUNCTION DEFINITION.
        # Copy this code into all model adapters published from this repo.
        # Dependencies should be added to /transforms-model-training/conda_recipe/meta.yaml
        from main._version import __version__ as generated_version_tag
        return PythonEnvironment(
            conda_dependencies=[
            CondaDependency(
                 "transforms-model-training-ri.stemma.main.repository.91e66421-692b-4338-84cb-27c6f1a1e785",
                 CondaVersionExact(version=f"{generated_version_tag}"),
                 "ri.stemma.main.repository.91e66421-692b-4338-84cb-27c6f1a1e785")
        ]
    )

模型推理转换代码 #

from transforms.api import transform, Input, Output
from palantir_models.transforms import ModelInput


@transform(
    model=ModelInput("/Modeling Tutorials/multi label classifier/multi_label_classifier"),
    inference_input=Input("/Modeling Tutorials/PubMed Multi Label Text Classification Dataset"),
    output=Output("/Modeling Tutorials/inference_output"),
)
def compute(inference_input, model, output):
    inference_results = model.transform(input_df=inference_input)   # 1. Call ModelAdapter.transform with the inputs specified in ModelAdapter.api
    df_out = inference_results.output_df                            # 2. Collect the desired output from the named tuple of inference result outputs
    output.write_pandas(df_out)

输出异常情况 #

推理转换仅返回22条结果（预期3000+条）：

问题排查与修复方案 #

1. 修正适配器输入列定义 #

适配器api()方法中，输入列名称和类型与实际训练、推理数据不匹配：原定义的features列不存在，且tuple类型错误。需修改为实际使用的特征列：

@classmethod
def api(cls):
    inputs = [
        ModelInput.Tabular(name="input_df",
                       df_type=DFType.PANDAS,
                       columns=[ModelApiColumn(name="abstractText", type=str)])
    ]
    outputs = [
        ModelOutput.Tabular(name="output_df",
                        columns=[ModelApiColumn(name="prediction", type=list)])
    ]
    return ModelApi(inputs, outputs)

2. 推理时提取正确的特征列 #

run_inference()方法中直接传入整个数据集给模型预测，但模型训练仅用了abstractText列，导致模型误将数据集列数当成样本数。需修改为提取目标列后再预测：

def run_inference(self, inputs, outputs):
    df_in = inputs.input_df
    df_out = outputs.output_df
    # 提取训练时使用的特征列
    X_in = df_in["abstractText"]
    predictions = self.model.predict(X_in)
    df = pd.DataFrame(predictions, columns=[self.prediction_column_name])
    df_out.write(df)

3. 修复模型保存方法的缩进错误 #

适配器save()方法中，pickle.dump(self.model, model_file)未缩进，会导致语法错误，模型无法正确保存。修正后：

def save(self, state_writer):
    with state_writer.open(ExampleModelAdapter.MODEL_SAVE_LOCATION, "wb") as model_file:
        pickle.dump(self.model, model_file)  # 增加缩进，确保在with代码块内执行
    with state_writer.open(ExampleModelAdapter.METADATA_SAVE_LOCATION, "wb") as metadata_file:
        metadata = {
            ExampleModelAdapter.PREDICTION_COLUMN_NAME_KEY: self.prediction_column_name
        }
        pickle.dump(metadata, metadata_file)

4. 验证推理输入数据格式 #

推理转换中需传入Pandas DataFrame而非Input对象，可添加验证代码确认输入行数：

def compute(inference_input, model, output):
    df = inference_input.pandas()
    print(f"Input dataset rows: {len(df)}")  # 打印输入行数，确认数据完整
    inference_results = model.transform(input_df=df)
    df_out = inference_results.output_df                            
    output.write_pandas(df_out)

内容的提问来源于stack exchange，提问作者tessa