如何修改Python标注代码实现多标签同时赋值到DataFrame对应列

我写了一段基于pandas的文本标注Python代码，输入单个标签（如label1）时能给DataFrame对应列赋值1，但输入多个标签（如label1, label2）时无法同时给多个对应列赋值1。希望修改代码实现该功能，以下是我的代码、当前输出及期望输出示例：

原代码 #

import pandas as pd
import time

data_pool = {'corpus': ['aa', 'bb', 'cc','dd', 'ee'], 'zero_level_name': ['a', 'b', 'c','d', 'e'], 'time': ['', '', '', '', ''], 'labels': ['', '', '', '', '']}

data_pool = pd.DataFrame(data_pool)

print(data_pool)

data_pool[['label1', 'label2', 'label3', 'label4']] = ''

index = 0
listA = ['label1', 'label2', 'label3', 'label4']
number_of_instances = int(input('Please, enter the number of texts you want to annotate today: '))

while index < number_of_instances:
    row = data_pool.loc[index]
    print("Enter labels for the text or enter 2 to go back to previous:")
    print()
    print('index:', index)
    print()
    print('text :\n\n\n', row['corpus'],'\n\n\n')
    start = time.time()
    label = input(": ")
    end = time.time()
    label.lower()
    if label == '2':
        index -= 1
        if index < 0:
            print('There is no previous row')
    else:
        label = label.split(',')
        label = [i.strip().lower() for i in label]
        for i in label:
            if i not in listA:
                print('Invalid input, try again')
                index -= 1
            else:
                if i == 'label1':
                    data_pool.loc[index, 'label1'] = 1
                elif i == 'label2':
                    data_pool.loc[index, 'label2'] = 1
                elif i == 'label3':
                    data_pool.loc[index, 'label3'] = 1
                elif i == 'label4':
                   data_pool.loc[index, 'label4'] = 1
                data_pool.loc[index, 'zero_level_name'] = label
                data_pool.loc[index, 'time'] = end-start
                break
        index += 1


print(data_pool)

当前输出 #

corpus zero_level_name      time labels label1 label2 label3 label4
0     aa        [label1]  5.372776             1                     
1     bb        [label2]  3.291902                    1               
2     cc               c                                             
3     dd               d                                             
4     ee               e   

期望输出（输入多个标签时） #

corpus zero_level_name      time labels label1 label2 label3 label4
0     aa  [label1, label2]  5.372776             1       1                     
1     bb  [label2, label3]  3.291902                    1     1         
2     cc  [label2, label3, label4] 3.548               1     1        1               
3     dd               d                                             
4     ee               e   

问题核心是原代码处理标签时，遍历到第一个有效标签就执行break，导致后续标签无法被处理。同时还有几个逻辑细节需要调整，修改后的完整代码如下：

import pandas as pd
import time

data_pool = {'corpus': ['aa', 'bb', 'cc','dd', 'ee'], 'zero_level_name': ['a', 'b', 'c','d', 'e'], 'time': ['', '', '', '', ''], 'labels': ['', '', '', '', '']}

data_pool = pd.DataFrame(data_pool)

print(data_pool)

data_pool[['label1', 'label2', 'label3', 'label4']] = ''

index = 0
listA = ['label1', 'label2', 'label3', 'label4']
number_of_instances = int(input('Please, enter the number of texts you want to annotate today: '))

while index < number_of_instances:
    # 每次处理当前行前，清空已有标签列，避免回退时残留旧值
    data_pool.loc[index, ['label1', 'label2', 'label3', 'label4']] = ''
    
    row = data_pool.loc[index]
    print("Enter labels for the text or enter 2 to go back to previous:")
    print()
    print('index:', index)
    print()
    print('text :\n\n\n', row['corpus'],'\n\n\n')
    start = time.time()
    label_input = input(": ")
    end = time.time()
    label_input = label_input.lower()  # 修正大小写处理：重新赋值给变量
    
    if label_input == '2':
        index -= 1
        if index < 0:
            print('There is no previous row')
            index = 0  # 避免index为负，卡住循环
    else:
        labels = label_input.split(',')
        labels = [i.strip() for i in labels]
        valid = True
        
        # 先批量校验所有标签是否合法
        for label in labels:
            if label not in listA:
                print('Invalid input, try again')
                valid = False
                break
        
        if valid:
            # 遍历所有有效标签，设置对应列值为1
            for label in labels:
                data_pool.loc[index, label] = 1
            
            # 更新元数据字段
            data_pool.loc[index, 'zero_level_name'] = labels
            data_pool.loc[index, 'time'] = end - start
            index += 1
        else:
            # 无效输入时不移动index，重新标注当前行
            continue

print(data_pool)

关键修改点 #

• 移除break语句：让循环能遍历所有输入的标签，完成多列赋值
• 提前校验所有标签：先检查所有输入标签的合法性，避免部分标签被错误写入
• 清空当前行标签列：每次进入循环时重置当前行的标签列，解决回退标注时旧值残留问题
• 修正大小写处理：label_input.lower()需要重新赋值给变量，否则原字符串不会改变

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