import torch.utils.data as data import pandas as pd import random from .config import * from .process import process_data from transformers import BertTokenizerFast #返回一个rel2id 和id2rel def get_rel(): df = pd.read_csv(REL_PATH, names=['rel', 'id']) return df['rel'].tolist(), dict(df.values) #生成长度为len，hot_pos位置为1其余位置为0的独热编码 def multihot(length, hot_pos): return [1 if i in hot_pos else 0 for i in range(length)] class Dataset(data.Dataset): def __init__(self, type='train'): super().__init__() _, self.rel2id = get_rel() # 加载文件 if type == 'train': file_path = TRAIN_PATH elif type == 'test': file_path = TEST_PATH self.lines = process_data(file_path) # 加载bert self.tokenizer = BertTokenizerFast.from_pretrained(BERT_MODEL_NAME) def __len__(self): return len(self.lines) def __getitem__(self, index): info = self.lines[index] # print(type(info)) # exit() tokenized = self.tokenizer(info['text'], return_offsets_mapping=True) info['input_ids'] = tokenized['input_ids'] info['offset_mapping'] = tokenized['offset_mapping'] return self.parse_dict(info) def get_pos_id(self, source, elem): for head_id in range(len(source)): tail_id = head_id + len(elem) if source[head_id:tail_id] == elem: return head_id, tail_id - 1 def collate_fn(self,batch): #获得最长的句子长度,便于填充pad batch.sort(key=lambda x : len(x["input_ids"]),reverse=True) max_len = len(batch[0]["input_ids"]) batch_text = { 'text': [], 'input_ids': [], 'offset_mapping': [], 'triple_list': [], } batch_mask = [] batch_sub = { 'heads_seq': [], 'tails_seq': [], } batch_sub_rnd = { 'head_seq': [], 'tail_seq': [], } batch_obj_rel = { 'heads_mx': [], 'tails_mx': [], } #对batch中的每一个item进行处理 for item in batch: input_ids = item["input_ids"] #对元素进行pad填充 item_len = len(input_ids) pad_len = max_len - item_len input_ids = input_ids + [0] * pad_len mask = [1] * item_len + [0] * pad_len # print(mask) # exit() sub_head_seq = multihot(max_len,item["sub_head_ids"]) sub_tail_seq = multihot(max_len, item["sub_tail_ids"]) # print(item["sub_head_ids"]) # print(sub_head_seq) # exit() if len(item['triple_id_list']) == 0: #如果没有三元组则continue continue sub_rnd = random.choice(item['triple_id_list'])[0] sub_rnd_head_seq = multihot(max_len, [sub_rnd[0]]) sub_rnd_tail_seq = multihot(max_len, [sub_rnd[1]]) #根据随机subject计算relations矩阵 obj_head_mx = [[0] * REL_SIZE for _ in range(max_len)] #生成两个二维全0矩阵（一个head矩阵一个tail矩阵） obj_tail_mx = [[0] * REL_SIZE for _ in range(max_len)] for triple in item["triple_id_list"]: #对全0矩阵进行填充获得obj的head和tail的rel矩阵 rel_id = triple[1] head_id, tail_id = triple[2] if triple[0] == sub_rnd: # 对于本课题可以取消这一步，因为数据集中有且仅有唯一的sub obj_head_mx[head_id][rel_id] = 1 obj_tail_mx[tail_id][rel_id] = 1 #重新组装batch，一条item压入一组信息 batch_text["text"].append(item["text"]) batch_text["input_ids"].append(input_ids) batch_text["offset_mapping"].append(item["offset_mapping"]) batch_text["triple_list"].append(item["triple_list"]) batch_mask.append(mask) batch_sub["heads_seq"].append(sub_head_seq) batch_sub["tails_seq"].append(sub_tail_seq) batch_sub_rnd["head_seq"].append(sub_rnd_head_seq) batch_sub_rnd["tail_seq"].append(sub_rnd_tail_seq) # print(sub_rnd_head_seq) # print(sub_head_seq) # exit() batch_obj_rel["heads_mx"].append(obj_head_mx) batch_obj_rel["tails_mx"].append(obj_tail_mx) return batch_mask,(batch_text,batch_sub_rnd),(batch_sub,batch_obj_rel) def parse_dict(self, info): #对dict串进行解析 text = info['text'] input_ids = info['input_ids'] #整个text的input_ids dct = { 'text': text, 'input_ids': input_ids, 'offset_mapping': info['offset_mapping'], 'sub_head_ids': [], 'sub_tail_ids': [], 'triple_list': [], 'triple_id_list': [] } for spo in info['spo_list']: subject = spo['subject'] object = spo['object'] predicate = spo['predicate'] dct['triple_list'].append((subject, predicate, object)) # 计算 subject 实体位置 tokenized = self.tokenizer(subject, add_special_tokens=False) sub_token = tokenized['input_ids'] sub_pos_id = self.get_pos_id(input_ids, sub_token) if not sub_pos_id: continue sub_head_id, sub_tail_id = sub_pos_id # 计算 object 实体位置 tokenized = self.tokenizer(object, add_special_tokens=False) obj_token = tokenized['input_ids'] obj_pos_id = self.get_pos_id(input_ids, obj_token) if not obj_pos_id: continue obj_head_id, obj_tail_id = obj_pos_id # 数据组装 dct['sub_head_ids'].append(sub_head_id) dct['sub_tail_ids'].append(sub_tail_id) dct['triple_id_list'].append(( [sub_head_id, sub_tail_id], self.rel2id[predicate], [obj_head_id, obj_tail_id], )) return dct # # 数据预处理部分已经格式化 # spo = info["spo_list"] # subject = spo['subject'] # object = spo['object'] # predicate = spo['predicate'] # dct['triple_list'].append((subject, predicate, object)) # # tokenized = self.tokenizer(subject, add_special_tokens=False) # sub_token = tokenized['input_ids'] # # print(input_ids) # # print(sub_token) # # exit() # sub_pos_id = self.get_pos_id(input_ids, sub_token) #通过编码后的 # # sub_head_id, sub_tail_id = sub_pos_id # # 计算 object 实体位置 # tokenized = self.tokenizer(object, add_special_tokens=False) # obj_token = tokenized['input_ids'] # obj_pos_id = self.get_pos_id(input_ids, obj_token) # # obj_head_id, obj_tail_id = obj_pos_id # # 数据组装 # dct['sub_head_ids'].append(sub_head_id) # dct['sub_tail_ids'].append(sub_tail_id) # # dct['triple_id_list'].append(( # [sub_head_id, sub_tail_id], # self.rel2id[predicate], # [obj_head_id, obj_tail_id], # )) # # return dct if __name__ == '__main__': dataset = Dataset() loader = data.DataLoader(dataset, shuffle=False, batch_size=2, collate_fn=dataset.collate_fn) print(next(iter(loader))) exit()