如何将TF-IDF数组的预测结果与原始Pandas DataFrame合并？

问题根源 #

你遇到的两个问题本质是直接对TF-IDF向量做train_test_split时，丢失了原DataFrame的行索引：

• Numpy数组（或稀疏矩阵）没有index()方法，这是Pandas数据结构独有的方法，所以调用会报错。
• array.indices是稀疏TF-IDF矩阵的内部属性，存储的是矩阵非零元素的位置索引，和原DataFrame的行索引完全无关，所以长度和测试集行数对不上是正常的。

两种解决方案 #

方案1：先拆分DataFrame，再做TF-IDF转换（推荐）

这是最稳妥的方式，拆分时直接保留原DataFrame的索引，后续直接把预测结果加回去就行：

from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

# 假设原DataFrame为df，文本列是'text'，目标列是'target'
# 先拆分DataFrame，X_train_df/X_test_df会保留原行索引
X_train_df, X_test_df, y_train, y_test = train_test_split(
    df[['text']], df['target'], test_size=0.2, random_state=42
)

# 对训练集拟合TF-IDF并转换，测试集直接转换
tfidf = TfidfVectorizer(max_features=5000)
X_train_tfidf = tfidf.fit_transform(X_train_df['text'])
X_test_tfidf = tfidf.transform(X_test_df['text'])

# 训练模型并预测
model = LogisticRegression()
model.fit(X_train_tfidf, y_train)
y_pred = model.predict(X_test_tfidf)

# 直接把预测结果合并到测试集DataFrame
X_test_df['prediction'] = y_pred

方案2：已拆分TF-IDF向量时，补回原索引

如果已经完成了TF-IDF转换和拆分，那可以在拆分时同时把原DataFrame的索引也拆分出来，再对应合并：

from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

# 假设已生成整个数据集的TF-IDF矩阵tfidf_matrix，原DataFrame为df
# 拆分时同时拆分特征矩阵、目标变量、原行索引
X_train_tfidf, X_test_tfidf, y_train, y_test, train_idx, test_idx = train_test_split(
    tfidf_matrix, df['target'], df.index, test_size=0.2, random_state=42
)

# 训练模型得到预测结果
model = LogisticRegression()
model.fit(X_train_tfidf, y_train)
y_pred = model.predict(X_test_tfidf)

# 根据test_idx从原DataFrame中取出测试集，再加入预测结果
X_test_df = df.loc[test_idx].copy()
X_test_df['prediction'] = y_pred

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