6分类任务中如何将句子与模型Logits结果转换为DataFrame？

核心思路：批量处理+一次性构造DataFrame #

循环单个处理句子并逐行添加到DataFrame的效率极低，尤其当数据量较大时。最优方式是批量处理所有句子，一次性生成完整的DataFrame，既减少模型调用开销，也避免了逐行操作的性能损耗。

完整代码实现 #

import pandas as pd
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "Emanuel/bertweet-emotion-base"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 目标句子列表
sentences = ['I like the weather today', 'The movie was very scary', 'Love is in the air']

# 批量Tokenize：开启padding和truncation保证输入长度统一
inputs = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True)
# 获取模型输出的Logits（传入所有tokenizer输出，包括attention_mask，避免计算错误）
outputs = model(**inputs)
# 转换为numpy数组（脱离计算图）
logits = outputs.logits.detach().numpy()

# 构造Logits的DataFrame，自动命名class_1到class_6
logits_df = pd.DataFrame(logits, columns=[f"class_{i+1}" for i in range(logits.shape[1])])
# 在最左侧插入sentence列
logits_df.insert(0, "sentence", sentences)

print(logits_df)

输出结果 #

sentence   class_1   class_2   class_3   class_4   class_5   class_6
0  I like the weather today -0.839088  2.948057 -0.513454  0.703865 -0.501967 -2.619496
1    The movie was very scary -0.884791 -0.964207 -2.210887 -0.439322  4.338617 -0.373839
2        Love is in the air -0.487504  3.294920  2.166052 -0.645325 -1.710200 -2.817954

替代方案：循环收集数据后一次性构造 #

如果因内存限制无法批量处理所有句子，可以先将每个句子的结果收集到列表中，最后再生成DataFrame（禁止逐行append到DataFrame，因为每次append都会创建新对象，性能极差）：

import pandas as pd
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "Emanuel/bertweet-emotion-base"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

sentences = ['I like the weather today', 'The movie was very scary', 'Love is in the air']
data_list = []

for sent in sentences:
    inputs = tokenizer(sent, return_tensors="pt")
    outputs = model(**inputs)
    # 将Logits从二维数组转为一维，拼接句子后加入列表
    logit_row = outputs.logits.detach().numpy().flatten()
    data_list.append([sent] + list(logit_row))

# 一次性构造DataFrame
df = pd.DataFrame(data_list, columns=["sentence"] + [f"class_{i+1}" for i in range(6)])

关键注意点 #

1. 批量处理时必须开启padding=True和truncation=True，确保所有输入序列长度一致，符合模型要求。
2. 传入模型时尽量用**inputs解包所有tokenizer输出（包括attention_mask），避免模型因忽略注意力掩码导致计算错误。
3. 避免逐行操作DataFrame，优先采用批量或列表收集后一次性构造的方式。

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