上一节课，我们重新整理了样本文件，接下来就可以定义Dataset类，来加载数据了。通过观察可以发现，样本数据中，每个句子只有一个实体。但在真实业务场景中，一个句子有多个实体，且有不同情感极性的情况也很常见。

为了让模型也可以适用于，一个句子有多个实体的情况，我们把样本中相邻两个句子合并，用最小的代价，模拟出多个实体的场景。

代码示例

1、添加配置项

# config.py
BIO_O_ID = 0
BIO_B_ID = 1
BIO_I_ID = 2
BIO_MAP = {'O': BIO_O_ID, 'B-ASP': BIO_B_ID, 'I-ASP': BIO_I_ID}
ENT_SIZE = 3

POLA_O_ID = -1
POLA_MAP = ['Negative', 'Positive']
POLA_DIM = 2

BERT_PAD_ID = 0
BERT_MODEL_NAME = 'bert-base-chinese'
BERT_DIM = 768

SRD = 3  # Semantic-Relative Distance

2、模块引入

# utils.py
import pandas as pd
import torch.utils.data as data
from config import *
from transformers import BertTokenizer
import torch
import random

3、定义Dataset类

class Dataset(data.Dataset):
    def __init__(self, type='train'):
        super().__init__()
        file_path = TRAIN_FILE_PATH if type == 'train' else TEST_FILE_PATH
        self.df = pd.read_csv(file_path)
        self.tokenizer = BertTokenizer.from_pretrained(BERT_MODEL_NAME)

    def __len__(self):
        return len(self.df) - 1

    def __getitem__(self, index):
        pass

4、加载单条数据

class Dataset(data.Dataset):
    def __getitem__(self, index):
        # 相邻两个句子拼接
        text1, bio1, pola1 = self.df.loc[index]
        text2, bio2, pola2 = self.df.loc[index + 1]
        text = text1 + ' ; ' + text2
        bio = bio1 + ' O ' + bio2
        pola = pola1 + ' -1 ' + pola2
        # 按自己的规则分词
        tokens = ['[CLS]'] + text.split(' ') + ['[SEP]']
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
        # BIO标签转id
        bio_arr = ['O'] + bio.split(' ') + ['O']
        bio_label = [BIO_MAP[l] for l in bio_arr]
        # 情感值转数字
        pola_arr = ['-1'] + pola.split(' ') + ['-1']
        pola_label = list(map(int, pola_arr))
        return input_ids, bio_label, pola_label

if __name__ == '__main__':
    dataset = Dataset()
    loader = data.DataLoader(dataset)
    print(iter(loader).next())