Python 将数据写入 csv 时的错误预测

Question

我训练了一个情绪分析模型，然后我在 Python 中创建了一个函数来基于这个模型预测文本。

以下是从文本中获取情绪（负面、正面、中性）的函数：

def get_sentiment(text):
    
    encoded_review = tokenizer.encode_plus(

          text,

          max_length=MAX_LEN,

          add_special_tokens=True,

          return_token_type_ids=False,

          pad_to_max_length=True,

          return_attention_mask=True,

          return_tensors='pt',

)


    input_ids = encoded_review['input_ids'].to(device)

    attention_mask = encoded_review['attention_mask'].to(device)

    output = model(input_ids, attention_mask)

    _, prediction = torch.max(output, dim=1)
    
    return class_names[prediction]

当我用一个文本测试这个 get_sentiment 函数时，它给出了正确的情绪：

get_sentiment('Miksi kukaan ei soittanut minulle??? Olen odottanut kauan!')

消极的

现在我有一个巨大的 csv，其中包含我想用这个 get_sentiment 函数预测的文本。 我想在 csv 中输入文本旁边的情绪。 我使用了以下代码：

from csv import writer
from csv import reader

with open('input.csv', 'r') as read_obj,\
    open('output.csv', 'w', newline='') as write_obj:
        csv_reader = reader(read_obj)
        csv_writer = writer(write_obj)
        
        for row in csv_reader:
            row.append(get_sentiment(row))
            csv_writer.writerow(row)

它在 output.csv 中附加了文本和情绪，但现在的情绪不正确：

Miksi kukaan ei soittanut minulle？？？ Olen odottanut kauan!,中立

我哪里做错了？ 请帮忙。

先谢谢了。

- 编辑 -

下面回答了解决方案。 我两个都试过了，都成功了！ 非常感谢 ：）

Answer 1

如果我没记错的话，你的get_sentiment函数应该用文本（ str ）提供，但你正在做

for row in csv_reader:
    row.append(get_sentiment(row))
    csv_writer.writerow(row)

因此提供整行，即tuple 。 如果第一列中存在文本，那么您应该按照以下方式抓取第一列

for row in csv_reader:
    row.append(get_sentiment(row[0]))
    csv_writer.writerow(row)

Answer 2

with open('input.csv', 'r') as read_obj,\
    open('output.csv', 'w', newline='') as write_obj:
        for row in read_obj:
            write_obj.writerow(row + "," + get_sentiment(row))

那样有用吗？