如何在Scikit-learn管道中组合LLM嵌入、TF-IDF和元数据

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# 模拟LLM嵌入生成器（通常您会使用预训练模型，如Sentence Transformers）
def generate_llm_embedding(texts):
    """模拟LLM嵌入生成。"""
    print(f"Generating mock embeddings for {len(texts)} texts...")
    # 为每个文本生成一个固定维度的随机向量
    embedding_dim = 32
    embeddings = np.random.rand(len(texts), embedding_dim)
    return embeddings

# 模拟数据
data = {
    'text': [
        "这是一个关于机器学习的好故事",
        "深度学习正在改变世界",
        "自然语言处理是一个热门领域",
        "如何训练一个高效的神经网络",
        "本文讨论了数据科学的未来",
        "Scikit-learn管道的强大功能"
    ],
    'metadata_feature_1': [10, 25, 5, 50, 30, 15],
    'metadata_feature_2': ['A', 'B', 'A', 'C', 'B', 'A'],
    'target': [0, 1, 0, 1, 1, 0]
}
df = pd.DataFrame(data)

# 准备特征和目标
X = df[['text', 'metadata_feature_1', 'metadata_feature_2']]
y = df['target']

# 拆分数据
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

print("训练数据样本数:", len(X_train))
print("测试数据样本数:", len(X_test))

# 文本特征（将由TF-IDF处理）
text_features = 'text'

# 元数据特征
numeric_features = ['metadata_feature_1']
categorical_features = ['metadata_feature_2']

# 用于元数据的预处理器
from sklearn.preprocessing import StandardScaler, OneHotEncoder

# 组合数值和类别元数据
metadata_transformer = Pipeline(steps=[
    ('onehot', OneHotEncoder(handle_unknown='ignore'))
])

# 完整的数据转换器
preprocessor = ColumnTransformer(
    transformers=[
        ('tfidf', TfidfVectorizer(), text_features),
        ('metadata_numeric', StandardScaler(), numeric_features),
        ('metadata_categorical', metadata_transformer, categorical_features)
    ], 
    remainder='passthrough' # 保留任何未指定的列（在本例中，没有）
)

# 假设我们已经计算了嵌入
X_train_text = X_train['text'].tolist()
X_test_text = X_test['text'].tolist()

# 生成嵌入
llm_embed_train = generate_llm_embedding(X_train_text)
llm_embed_test = generate_llm_embedding(X_test_text)

print(f"训练嵌入形状: {llm_embed_train.shape}")

from sklearn.base import BaseEstimator, TransformerMixin

# 自定义转换器，用于合并预计算的嵌入和管道输出
class FeatureCombiner(BaseEstimator, TransformerMixin):
    def __init__(self, combiner_transformer, llm_embeddings):
        self.combiner_transformer = combiner_transformer
        self.llm_embeddings = llm_embeddings

    def fit(self, X, y=None):
        # 仅对组合器进行拟合
        self.combiner_transformer.fit(X)
        return self

    def transform(self, X):
        # 转换非嵌入特征
        non_llm_features = self.combiner_transformer.transform(X)

        # 确保输出是密集矩阵，以便堆叠
        if hasattr(non_llm_features, 'toarray'):
            non_llm_features = non_llm_features.toarray()

        # 水平堆叠
        combined_features = np.hstack((non_llm_features, self.llm_embeddings))
        print(f"Combined features shape: {combined_features.shape}")
        return combined_features

# 在fit阶段，我们需要将X_train传递给组合器进行拟合
# 但FeatureCombiner的transform阶段需要X_test（或X_train本身）

# 准备用于管道的数据（仅包含需要ColumnTransformer处理的列）
X_train_pipe = X_train
X_test_pipe = X_test

# 最终管道步骤：首先使用ColumnTransformer处理TF-IDF和元数据
# 然后使用FeatureCombiner合并LLM嵌入

# 1. 创建ColumnTransformer处理所有非LLM特征
text_processor = TfidfVectorizer(max_features=100) # 限制特征数量以保持示例可控

metadata_transformer_final = ColumnTransformer(
    transformers=[
        ('tfidf', text_processor, 'text'),
        ('numeric', StandardScaler(), numeric_features),
        ('categorical', OneHotEncoder(handle_unknown='ignore'), categorical_features)
    ], 
    remainder='drop' 
)

# 2. 创建合并器，它将使用ColumnTransformer进行拟合/转换，并附加预先计算的嵌入
# 注意：这里我们使用一个特殊技巧，因为LLM嵌入是在外部计算的，我们不能直接在管道内拟合它。
# 我们需要确保FeatureCombiner的fit只拟合ColumnTransformer。

# 为了满足Scikit-learn API，我们需要一个更简单的设计：
# 1. 拟合ColumnTransformer（TF-IDF/元数据）
metadata_transformer_final.fit(X_train_pipe)

# 2. 转换所有数据
train_metadata_features = metadata_transformer_final.transform(X_train_pipe)
test_metadata_features = metadata_transformer_final.transform(X_test_pipe)

# 确保都是密集矩阵
if hasattr(train_metadata_features, 'toarray'):
    train_metadata_features = train_metadata_features.toarray()
if hasattr(test_metadata_features, 'toarray'):
    test_metadata_features = test_metadata_features.toarray()

print(f"TF-IDF/元数据训练特征形状: {train_metadata_features.shape}")

# 3. 水平堆叠（Hstack）
# 训练集组合
X_train_combined = np.hstack((train_metadata_features, llm_embed_train))
# 测试集组合
X_test_combined = np.hstack((test_metadata_features, llm_embed_test))

print(f"最终训练特征形状: {X_train_combined.shape}")

# 4. 构建最终分类器管道
# 现在我们有了一个包含所有特征的密集矩阵，可以直接用于分类器
classifier_pipeline = Pipeline(steps=[
    ('classifier', LogisticRegression(max_iter=1000, solver='liblinear'))
])

# 训练模型
print("\n开始训练最终模型...")
classifier_pipeline.fit(X_train_combined, y_train)

# 预测
y_pred = classifier_pipeline.predict(X_test_combined)

# 评估
accuracy = accuracy_score(y_test, y_pred)

print(f"\n模型准确率 (组合特征): {accuracy:.4f}")

# ----------------------------------------------------------------------
# 理论上，如果使用一个统一的管道，它将如下所示（但需要自定义转换器来处理外部计算的嵌入）
# 实际项目中，分离预计算和管道训练是更常见的做法。
# ----------------------------------------------------------------------

# 示例：仅使用ColumnTransformer（不包含LLM嵌入，因为它们不是数据列的一部分）
# full_pipeline = Pipeline(steps=[
#     ('preprocessor', preprocessor),
#     ('classifier', LogisticRegression())
# ])
# full_pipeline.fit(X_train, y_train)