手把手医学知识图谱搭建案例
注：大家觉得博客好的话，别忘了点赞收藏呀，本人每周都会更新关于人工智能和大数据相关的内容，内容多为原创，Python Java Scala SQL 代码，CV NLP 推荐系统等，Spark Flink Kafka Hbase Hive Flume等等~写的都是纯干货，各种顶会的论文解读，一起进步。
今天和大家分享一下医学知识图谱中三元组搭建的案例
github: /king-yyf/CMeKG_tools
#博学谷IT学习技术支持#
文章目录
手把手医学知识图谱搭建案例前言一、先来看最终效果展示二、核心代码1.引入库2.训练数据3.数据预处理4.平平无奇的ner模型Model4s5.很有创意的Model4po6.正常训练模型7.demo案例总结前言
知识图谱（Knowledge Graph）被运用在很多科研领域，其重要程度不言而喻，很多大厂都在致力于搭建属于自己的知识图谱，尤其是知识抽取是知识图谱的核心，今天和大家分享一下医学知识图谱中三元组搭建的案例。
一、先来看最终效果展示
搭建出来的三元组准确性还是相当不错的
再来看看最终的知识图谱效果图,通过模型有了三元组，利用Neo4j搭建知识图谱so easy~
二、核心代码
1.引入库
这里最重要是导入transformers包，因为整个预训练模型是通过huggingface去做的。还是是用bert。
代码如下（示例）：
import jsonimport numpy as npimport torchimport torch.nn as nnfrom transformers import BertTokenizer, BertModel, AdamWfrom itertools import cycleimport gcimport randomimport timeimport re#先指定好所有的路径class config:batch_size = 1max_seq_len = 256num_p = 23learning_rate = 1e-5EPOCH = 2PATH_SCHEMA = "/CMeKG/predicate.json"PATH_TRAIN = '/CMeKG/train_example.json'PATH_BERT = "/CMeKG/model/medical_re"PATH_MODEL = "/CMeKG/model/medical_re/model_re.pkl"PATH_SAVE = '/CMeKG/model/save'tokenizer = BertTokenizer.from_pretrained("/CMeKG/model/medical_re/vocab.txt")id2predicate = {}predicate2id = {}
2.训练数据
其中PATH_TRAIN = '/CMeKG/train_example.json’是训练数据
{
“text”: “12小时尿沉渣计数的相关疾病:单纯型尿路感染，妊娠合并急性膀胱炎，慢性肾炎，狼疮性肾炎，急性膀胱炎12小时尿沉渣计数的相关症状是高血压，男子性功能障碍，蛋白尿，血尿，水肿，排尿困难及尿潴留，尿频伴尿急和尿痛”,
“spo_list”: [
[
“12小时尿沉渣计数”,
“相关疾病”,
“单纯型尿路感染”
]
其中text是原文，spo是三元组，s代表主体，p代表关系，s代表客体。
3.数据预处理
这里我都已经写好注视了，这个一个数据预处理逻辑
代码如下（示例）：
class IterableDataset(torch.utils.data.IterableDataset):def __init__(self, data, random):super(IterableDataset).__init__()self.data = dataself.random = randomself.tokenizer = config.tokenizerdef __len__(self):return len(self.data)def search(self, sequence, pattern):n = len(pattern)for i in range(len(sequence)):if sequence[i:i + n] == pattern:return ireturn -1def process_data(self):idxs = list(range(len(self.data)))if self.random:np.random.shuffle(idxs)batch_size = config.batch_sizemax_seq_len = config.max_seq_lennum_p = config.num_pbatch_token_ids = np.zeros((batch_size, max_seq_len), dtype=np.int)batch_mask_ids = np.zeros((batch_size, max_seq_len), dtype=np.int)batch_segment_ids = np.zeros((batch_size, max_seq_len), dtype=np.int)batch_subject_ids = np.zeros((batch_size, 2), dtype=np.int)batch_subject_labels = np.zeros((batch_size, max_seq_len, 2), dtype=np.int)batch_object_labels = np.zeros((batch_size, max_seq_len, num_p, 2), dtype=np.int)batch_i = 0for i in idxs:text = self.data[i]['text']batch_token_ids[batch_i, :] = self.tokenizer.encode(text, max_length=max_seq_len, pad_to_max_length=True,add_special_tokens=True)batch_mask_ids[batch_i, :len(text) + 2] = 1#对pad出来的设置成0spo_list = self.data[i]['spo_list']idx = np.random.randint(0, len(spo_list), size=1)[0]#相当于每次都是随机选一个S来组成数据s_rand = self.tokenizer.encode(spo_list[idx][0])[1:-1]#S的ID编码s_rand_idx = self.search(list(batch_token_ids[batch_i, :]), s_rand)#S所在text的开始索引位置batch_subject_ids[batch_i, :] = [s_rand_idx, s_rand_idx + len(s_rand) - 1]#S所在text的起始和终止索引位置for i in range(len(spo_list)):spo = spo_list[i]s = self.tokenizer.encode(spo[0])[1:-1]#不要首尾特殊字符p = config.prediction2id[spo[1]]o = self.tokenizer.encode(spo[2])[1:-1]s_idx = self.search(list(batch_token_ids[batch_i]), s)#S的开始位置o_idx = self.search(list(batch_token_ids[batch_i]), o)#O的开始位置if s_idx != -1 and o_idx != -1:#他俩都存在的话batch_subject_labels[batch_i, s_idx, 0] = 1#到时候要预测每一个token是不是S的起始和终止位置batch_subject_labels[batch_i, s_idx + len(s) - 1, 1] = 1if s_idx == s_rand_idx:batch_object_labels[batch_i, o_idx, p, 0] = 1#记录O的开始位置及S与O之间的关系batch_object_labels[batch_i, o_idx + len(o) - 1, p, 1] = 1#记录O的结束位置及S与O之间的关系batch_i += 1if batch_i == batch_size or i == idxs[-1]:yield batch_token_ids, batch_mask_ids, batch_segment_ids, batch_subject_labels, batch_subject_ids, batch_object_labelsbatch_token_ids[:, :] = 0batch_mask_ids[:, :] = 0batch_subject_ids[:, :] = 0batch_subject_labels[:, :, :] = 0batch_object_labels[:, :, :, :] = 0batch_i = 0def get_stream(self):return cycle(self.process_data())def __iter__(self):return self.get_stream()
4.平平无奇的ner模型Model4s
这里和普通的ner任务完全一样，就是调用huggingface的bert预训练模型接口
代码如下（示例）：
class Model4s(nn.Module):def __init__(self, hidden_size=768):super(Model4s, self).__init__()self.bert = BertModel.from_pretrained(config.PATH_BERT)for param in self.bert.parameters():param.requires_grad = Trueself.dropout = nn.Dropout(p=0.2)self.linear = nn.Linear(in_features=hidden_size, out_features=2, bias=True)self.sigmoid = nn.Sigmoid()def forward(self, input_ids, input_mask, segment_ids, hidden_size=768):hidden_states = self.bert(input_ids,attention_mask=input_mask,token_type_ids=segment_ids)[0] # (batch_size, sequence_length, hidden_size)output = self.sigmoid(self.linear(self.dropout(hidden_states))).pow(2)return output, hidden_states
5.很有创意的Model4po
这里通过上面的Model4s找到text主题后，固定主题，然后找到相应的客体和关系，相当精彩的想法。也是整个项目精彩之处。
代码如下（示例）：
class Model4po(nn.Module):def __init__(self, num_p=config.num_p, hidden_size=768):super(Model4po, self).__init__()self.dropout = nn.Dropout(p=0.4)self.linear = nn.Linear(in_features=hidden_size, out_features=num_p * 2, bias=True)self.sigmoid = nn.Sigmoid()def forward(self, hidden_states, batch_subject_ids, input_mask):all_s = torch.zeros((hidden_states.shape[0], hidden_states.shape[1], hidden_states.shape[2]),dtype=torch.float32)for b in range(hidden_states.shape[0]):s_start = batch_subject_ids[b][0]s_end = batch_subject_ids[b][1]s = hidden_states[b][s_start] + hidden_states[b][s_end]#起始特征+终止特征cue_len = torch.sum(input_mask[b])#实际长度all_s[b, :cue_len, :] = s#将所有位置的特征设置成主体的hidden_states += all_s#每一个位置实际特征都是 自身 + 主体#我估计pow(4)这么大个数 是由于预测出来的结果都有些大，要降低预测值大小output = self.sigmoid(self.linear(self.dropout(hidden_states))).pow(4)#预测每一个位置与主题的关系return output # (batch_size, max_seq_len, num_p*2)
6.正常训练模型
写的也比较普通，把2个模型的损失加起来就行。
def train(train_data_loader, model4s, model4po, optimizer):for epoch in range(config.EPOCH):begin_time = time.time()model4s.train()model4po.train()train_loss = 0.for bi, batch in enumerate(train_data_loader):if bi >= len(train_data_loader) // config.batch_size:breakbatch_token_ids, batch_mask_ids, batch_segment_ids, batch_subject_labels, batch_subject_ids, batch_object_labels = batchbatch_token_ids = torch.tensor(batch_token_ids, dtype=torch.long)batch_mask_ids = torch.tensor(batch_mask_ids, dtype=torch.long)batch_segment_ids = torch.tensor(batch_segment_ids, dtype=torch.long)batch_subject_labels = torch.tensor(batch_subject_labels, dtype=torch.float)batch_object_labels = torch.tensor(batch_object_labels, dtype=torch.float).view(config.batch_size,config.max_seq_len,config.num_p * 2)batch_subject_ids = torch.tensor(batch_subject_ids, dtype=torch.int)batch_subject_labels_pred, hidden_states = model4s(batch_token_ids, batch_mask_ids, batch_segment_ids)loss4s = loss_fn(batch_subject_labels_pred, batch_subject_labels.to(torch.float32))loss4s = torch.mean(loss4s, dim=2, keepdim=False) * batch_mask_ids#只计算非pad部分loss4s = torch.sum(loss4s)loss4s = loss4s / torch.sum(batch_mask_ids)batch_object_labels_pred = model4po(hidden_states, batch_subject_ids, batch_mask_ids)loss4po = loss_fn(batch_object_labels_pred, batch_object_labels.to(torch.float32))loss4po = torch.mean(loss4po, dim=2, keepdim=False) * batch_mask_idsloss4po = torch.sum(loss4po)loss4po = loss4po / torch.sum(batch_mask_ids)loss = loss4s + loss4pooptimizer.zero_grad()loss.backward()optimizer.step()train_loss += float(loss.item())print('batch:', bi, 'loss:', float(loss.item()))print('final train_loss:', train_loss / len(train_data_loader) * config.batch_size, 'cost time:',time.time() - begin_time)del train_data_loadergc.collect();return {"model4s_state_dict": model4s.state_dict(),"model4po_state_dict": model4po.state_dict(),"optimizer_state_dict": optimizer.state_dict(),}
7.demo案例
输出的结果就是文章开头
import medical_reimport jsonmodel4s, model4po = medical_re.load_model()text = '据报道称，新冠肺炎患者经常会发热、咳嗽，少部分患者会胸闷、乏力，其病因包括: 1.自身免疫系统缺陷\n2.人传人。' # content是输入的一段文字res = medical_re.get_triples(text, model4s, model4po)print(json.dumps(res, ensure_ascii=False, indent=True))
总结
通过huggingface中bert预训练模型实现的一个非常不错的知识图谱搭建过程。Neo4j怎么搭建，以后有时间继续更新。
如果觉得《手把手医学知识图谱搭建案例》对你有帮助，请点赞、收藏，并留下你的观点哦！
