🧮 第4课：知识表示学习

import numpy as np from collections import defaultdict class TransE: """TransE知识表示学习模型""" def __init__(self, entities, relations, dim=50, margin=1.0, lr=0.01): self.entities = list(entities) self.relations = list(relations) self.dim = dim self.margin = margin self.lr = lr # 初始化嵌入向量（Xavier初始化） scale = np.sqrt(6.0 / (dim + dim)) self.entity_embeddings = { e: np.random.uniform(-scale, scale, dim) for e in self.entities } self.relation_embeddings = { r: np.random.uniform(-scale, scale, dim) for r in self.relations } # 归一化 self._normalize() def _normalize(self): """归一化实体向量到单位球面""" for e in self.entity_embeddings: norm = np.linalg.norm(self.entity_embeddings[e]) if norm > 0: self.entity_embeddings[e] /= norm def distance(self, h, r, t): """计算三元组的距离（L1范数）""" h_vec = self.entity_embeddings[h] r_vec = self.relation_embeddings[r] t_vec = self.entity_embeddings[t] return np.linalg.norm(h_vec + r_vec - t_vec, ord=1) def _generate_negative(self, h, r, t): ">>>生成负样本：随机替换头实体或尾实体""" if np.random.random() < 0.5: h_neg = np.random.choice(self.entities) while h_neg == h: h_neg = np.random.choice(self.entities) return (h_neg, r, t) else: t_neg = np.random.choice(self.entities) while t_neg == t: t_neg = np.random.choice(self.entities) return (h, r, t_neg) def train(self, triples, epochs=100): """训练TransE模型""" for epoch in range(epochs): total_loss = 0 np.random.shuffle(triples) for h, r, t in triples: # 正样本距离 d_pos = self.distance(h, r, t) # 负样本 h_neg, r_neg, t_neg = self._generate_negative(h, r, t) d_neg = self.distance(h_neg, r_neg, t_neg) # Margin损失 loss = max(0, self.margin + d_pos - d_neg) total_loss += loss # 梯度更新（简化版SGD） if loss > 0: h_vec = self.entity_embeddings[h] r_vec = self.relation_embeddings[r] t_vec = self.entity_embeddings[t] diff = h_vec + r_vec - t_vec sign = np.sign(diff) self.entity_embeddings[h] -= self.lr * sign self.relation_embeddings[r] -= self.lr * sign self.entity_embeddings[t] += self.lr * sign # 负样本梯度 self.entity_embeddings[h_neg] += self.lr * sign * 0.5 self.entity_embeddings[t_neg] -= self.lr * sign * 0.5 self._normalize() if (epoch + 1) % 20 == 0: print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss:.4f}") def predict_tail(self, h, r, top_k=5): """预测最可能的尾实体（链接预测）""" scores = [] for e in self.entities: d = self.distance(h, r, e) scores.append((e, d)) scores.sort(key=lambda x: x[1]) return scores[:top_k] def similarity(self, e1, e2): """计算两实体的余弦相似度""" v1 = self.entity_embeddings[e1] v2 = self.entity_embeddings[e2] cos = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-8) return cos # ========== 训练示例 ========== np.random.seed(42) entities = {"鲁迅", "老舍", "徐志摩", "呐喊", "彷徨", "骆驼祥子", "再别康桥", "绍兴", "北京", "海宁"} relations = {"创作", "出生地"} triples = [ ("鲁迅", "创作", "呐喊"), ("鲁迅", "创作", "彷徨"), ("鲁迅", "出生地", "绍兴"), ("老舍", "创作", "骆驼祥子"), ("老舍", "出生地", "北京"), ("徐志摩", "创作", "再别康桥"), ("徐志摩", "出生地", "海宁"), ] model = TransE(entities, relations, dim=20, margin=1.0, lr=0.05) model.train(triples, epochs=100) # 链接预测 print(" === 链接预测: (鲁迅, 创作, ?) ===") for entity, dist in model.predict_tail("鲁迅", "创作"): print(f" {entity}: 距离={dist:.4f}") print(" === 链接预测: (老舍, 出生地, ?) ===") for entity, dist in model.predict_tail("老舍", "出生地"): print(f" {entity}: 距离={dist:.4f}") # 实体相似度 print(" === 实体相似度 ===") pairs = [("鲁迅", "老舍"), ("鲁迅", "徐志摩"), ("鲁迅", "呐喊"), ("呐喊", "彷徨")] for e1, e2 in pairs: sim = model.similarity(e1, e2) print(f" sim({e1}, {e2}) = {sim:.4f}")

# TransH模型实现 class TransH: """TransH: 在关系特定的超平面上做翻译""" def __init__(self, entities, relations, dim=50, margin=1.0): self.dim = dim self.margin = margin scale = np.sqrt(6.0 / (dim + dim)) self.entity_embeddings = {e: np.random.uniform(-scale, scale, dim) for e in entities} self.relation_embeddings = {r: np.random.uniform(-scale, scale, dim) for r in relations} self.normal_vectors = {r: np.random.uniform(-scale, scale, dim) for r in relations} self._normalize() def _normalize(self): for e in self.entity_embeddings: self.entity_embeddings[e] /= (np.linalg.norm(self.entity_embeddings[e]) + 1e-8) for r in self.normal_vectors: self.normal_vectors[r] /= (np.linalg.norm(self.normal_vectors[r]) + 1e-8) def project(self, entity, relation): """将实体投影到关系超平面""" e_vec = self.entity_embeddings[entity] n_vec = self.normal_vectors[relation] return e_vec - np.dot(e_vec, n_vec) * n_vec def distance(self, h, r, t): h_proj = self.project(h, r) t_proj = self.project(t, r) r_vec = self.relation_embeddings[r] return np.linalg.norm(h_proj + r_vec - t_proj, ord=1) def predict_tail(self, h, r, entities, top_k=5): scores = [(e, self.distance(h, r, e)) for e in entities] scores.sort(key=lambda x: x[1]) return scores[:top_k] # 测试TransH np.random.seed(42) model_h = TransH(entities, relations, dim=20) print("=== TransH链接预测 ===") for entity, dist in model_h.predict_tail("鲁迅", "创作", list(entities)): print(f" {entity}: 距离={dist:.4f}")

def evaluate_model(model, test_triples, all_entities): """评估知识表示学习模型""" ranks = [] for h, r, t in test_triples: # 预测尾实体 scores = [(e, model.distance(h, r, e)) for e in all_entities] scores.sort(key=lambda x: x[1]) rank = next(i + 1 for i, (e, _) in enumerate(scores) if e == t) ranks.append(rank) mr = np.mean(ranks) mrr = np.mean([1.0 / r for r in ranks]) hits1 = np.mean([1 if r <= 1 else 0 for r in ranks]) hits3 = np.mean([1 if r <= 3 else 0 for r in ranks]) hits10 = np.mean([1 if r <= 10 else 0 for r in ranks]) return {"MR": mr, "MRR": mrr, "Hits@1": hits1, "Hits@3": hits3, "Hits@10": hits10} print("=== TransE模型评估 ===") metrics = evaluate_model(model, triples, list(entities)) for k, v in metrics.items(): print(f" {k}: {v:.4f}")