1
# 知识图谱中的结构化知识示例
2
knowledge_graph = {
3
    ("Obama", "born_in", "Honolulu"): True,
4
    ("Obama", "nationality", "American"): True,
5
    ("Honolulu", "located_in", "Hawaii"): True,
6
    ("Hawaii", "part_of", "USA"): True,
7
}
8

9
# ERNIE 目标：让模型理解这些关系
10
# 普通 BERT：可能只是学会了语法模式
11
# ERNIE：理解实体和关系

1
class ERNIE(nn.Module):
2
    """ERNIE 模型核心组件"""
3

4
    def __init__(self, vocab_size, d_model, n_layers, entity_dim):
5
        super().__init__()
6

7
        # 文本编码器（类似 BERT）
8
        self.token_embedding = nn.Embedding(vocab_size, d_model)
9
        self.position_embedding = nn.Embedding(512, d_model)
10
        self.text_encoder = TransformerEncoder(n_layers, d_model)
11

12
        # 实体编码器
13
        self.entity_embedding = nn.Embedding(num_entities, entity_dim)
14
        self.knowledge_encoder = KnowledgeTransformer(n_layers, d_model, entity_dim)
15

16
        # 融合层
17
        self.fusion_layer = nn.Linear(d_model + entity_dim, d_model)
18
        self.ln_f = nn.LayerNorm(d_model)
19

20
    def forward(self, token_ids, entity_ids, token_to_entity_map):
21
        # 文本编码
22
        token_emb = self.token_embedding(token_ids)
23
        text_h = self.text_encoder(token_emb)
24

25
        # 实体编码
26
        entity_emb = self.entity_embedding(entity_ids)
27
        knowledge_h = self.knowledge_encoder(entity_emb, token_to_entity_map)
28

29
        # 融合
30
        fused = torch.cat([text_h, knowledge_h], dim=-1)
31
        output = self.ln_f(self.fusion_layer(fused))
32

33
        return output

1
class KnowledgeTransformer(nn.Module):
2
    """知识编码器：处理实体级别信息"""
3

4
    def __init__(self, d_model, entity_dim, n_heads):
5
        super().__init__()
6
        self.entity_proj = nn.Linear(entity_dim, d_model)
7
        self.attn = MultiHeadAttention(d_model, n_heads)
8
        self.ffn = PositionwiseFFN(d_model)
9

10
    def forward(self, entity_emb, token_to_entity_map, token_mask):
11
        # 将实体嵌入映射到文本空间
12
        entity_hidden = self.entity_proj(entity_emb)
13

14
        # 实体- token 对齐注意力
15
        # 对于每个 token，聚合相关实体的信息
16
        output = self.attn(
17
            query=token_hidden,
18
            key=entity_hidden,
19
            value=entity_hidden,
20
            mask=token_mask
21
        )
22

23
        return output + token_hidden

1
# ERNIE 的预训练任务
2

3
# 1. 遮蔽语言模型（标准 BERT）
4
def masked_language_model_loss(tokens, masked_tokens, predictions):
5
    """MLM 损失"""
6
    loss = F.cross_entropy(predictions, masked_tokens)
7
    return loss
8

9
# 2. 遮蔽实体预测（ERNIE 新增）
10
def masked_entity_model_loss(hidden, entity_labels):
11
    """预测被遮蔽的实体"""
12
    entity_logits = self.entity_head(hidden)
13
    loss = F.cross_entropy(entity_logits, entity_labels)
14
    return loss
15

16
# 3. 关系预测（ERNIE 新增）
17
def relation_prediction_loss(hidden, relations):
18
    """预测实体之间的关系"""
19
    # 使用实体表示计算关系分数
20
    entity_hidden = gather_entities(hidden, entity_positions)
21
    rel_logits = self.relation_head(entity_hidden)
22
    loss = F.cross_entropy(rel_logits, relations)
23
    return loss
24

25
# 总损失
26
total_loss = mlm_loss + mem_loss + rel_loss

1
# 测试模型是否学习了知识
2

3
# 原始句子
4
text = "[E1]Obama[/E1] was born in [E2]Honolulu[/E2]"
5

6
# ERNIE 能正确回答：
7
# Q: Where was Obama born?
8
# A: Honolulu (通过知识图谱增强理解)
9

10
# 测试关系推理
11
text = "[E1]Microsoft[/E1] was founded by [E2]Bill Gates[/E2]"
12

13
# ERNIE 能理解：
14
# - Microsoft 和 Bill Gates 的关系
15
# - 实体类型信息