基于知识图谱嵌入的自动化问答生成

import torch
import torch.nn as nn
import numpy as np

# 定义TransE模型
class TransE(nn.Module):
    def __init__(self, num_entities, num_relations, embedding_dim, margin=1.0):
        super(TransE, self).__init__()
        self.entity_embeddings = nn.Embedding(num_entities, embedding_dim)
        self.relation_embeddings = nn.Embedding(num_relations, embedding_dim)
        self.margin = margin
        self.loss_fn = nn.MarginRankingLoss(margin=margin)
        
        # 初始化嵌入矩阵
        nn.init.xavier_uniform_(self.entity_embeddings.weight)
        nn.init.xavier_uniform_(self.relation_embeddings.weight)
    
    def forward(self, head, relation, tail):
        head_emb = self.entity_embeddings(head)
        relation_emb = self.relation_embeddings(relation)
        tail_emb = self.entity_embeddings(tail)
        score = torch.norm(head_emb + relation_emb - tail_emb, p=1, dim=1)
        return score

    def compute_loss(self, pos_score, neg_score):
        target = -torch.ones(pos_score.size())
        return self.loss_fn(pos_score, neg_score, target)

# 定义知识图谱数据
data = {
    "entities": ["高血压", "血管收缩", "药物A", "头痛", "光合作用", "植物", "阳光"],
    "relations": ["由...引发", "治疗药物", "可能副作用", "发生在", "需要...进行"],
    "triples": [
        (0, 0, 1),  # 高血压 由...引发 血管收缩
        (0, 1, 2),  # 高血压 治疗药物 药物A
        (2, 2, 3),  # 药物A 可能副作用 头痛
        (4, 3, 5),  # 光合作用 发生在 植物
        (5, 4, 6)   # 植物 需要...进行 阳光
    ]
}

# 实体和关系数
num_entities = len(data["entities"])
num_relations = len(data["relations"])
embedding_dim = 100  # 嵌入维度

# 创建模型
model = TransE(num_entities, num_relations, embedding_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 模型训练
def train(model, data, epochs=100):
    for epoch in range(epochs):
        total_loss = 0
        for triple in data["triples"]:
            head, relation, tail = torch.LongTensor([triple[0]]), torch.LongTensor([triple[1]]), torch.LongTensor([triple[2]])
            pos_score = model(head, relation, tail)

            # 负采样
            neg_tail = torch.LongTensor([np.random.randint(0, num_entities)])
            neg_score = model(head, relation, neg_tail)

            # 计算损失并优化
            loss = model.compute_loss(pos_score, neg_score)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
           

 total_loss += loss.item()

        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch + 1}/{epochs}, Loss: {total_loss}")

# 开始训练
train(model, data)

  question = "高血压有哪些副作用？"
  entities = ner_model.predict(question)
  print(entities)  # 输出 ['高血压']

  question = "高血压有哪些副作用？"
  relation = re_model.predict(question, entities)
  print(relation)  # 输出 ['副作用']

entity = '高血压'
entity_embedding = entity_embedding_matrix[entity_id]
print(entity_embedding)  # 输出 [0.23, 0.45, -0.11, ...] 一个低维的向量表示

relation = '副作用'
relation_embedding = relation_embedding_matrix[relation_id]
print(relation_embedding)  # 输出 [0.12, -0.34, 0.09, ...] 一个低维的向量表示

question_embedding = entity_embedding + relation_embedding
print(question_embedding)  # 输出问题的嵌入向量

  import numpy as np
  def cosine_similarity(vec1, vec2):
      return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
  
  # 假设候选答案的嵌入是 answer_embeddings
  similarities = [cosine_similarity(question_embedding, ans) for ans in answer_embeddings]
  best_answer_index = np.argmax(similarities)
  best_answer = answers[best_answer_index]
  print(f"最佳答案是: {best_answer}")

print(f"高血压的副作用是: {best_answer}")