deepwalk.py

1"""2DeepWalk 图嵌入召回（第10章 §10.2.1）34流程：51. 在用户-物品二部图上做随机游走生成序列（Random Walk）62. 用 Skip-Gram（Word2Vec 思想）训练 Embedding73. 线上用 Item Embedding 做最近邻召回89此处用纯 Python 实现，不依赖 gensim/networkx，便于原理演示。10"""1112import random13import math14from collections import defaultdict151617# ── 1. 构建图 ─────────────────────────────────────────────────────────────────1819def build_graph(user_items: dict[str, list[str]]) -> dict[str, list[str]]:20    """21    构建无向二部图（用户节点 ↔ 物品节点）。22    返回邻接表 {node: [neighbor, ...]}。23    """24    graph: dict[str, list[str]] = defaultdict(list)25    for user, items in user_items.items():26        uid = f"u:{user}"27        for item in items:28            iid = f"i:{item}"29            graph[uid].append(iid)30            graph[iid].append(uid)31    return dict(graph)323334# ── 2. 随机游走 ───────────────────────────────────────────────────────────────3536def random_walk(graph: dict[str, list[str]], start: str, walk_len: int) -> list[str]:37    """从 start 节点出发进行随机游走，返回节点序列。"""38    walk = [start]39    for _ in range(walk_len - 1):40        neighbors = graph.get(walk[-1], [])41        if not neighbors:42            break43        walk.append(random.choice(neighbors))44    return walk454647def generate_walks(48    graph: dict[str, list[str]],49    num_walks: int = 10,50    walk_len: int = 20,51) -> list[list[str]]:52    """对图中所有节点各做 num_walks 次随机游走。"""53    nodes = list(graph.keys())54    walks = []55    for _ in range(num_walks):56        random.shuffle(nodes)57        for node in nodes:58            walks.append(random_walk(graph, node, walk_len))59    return walks606162# ── 3. Skip-Gram 训练（负采样简化版）────────────────────────────────────────6364class SkipGram:65    def __init__(self, vocab: list[str], dim: int = 32, lr: float = 0.025):66        self.vocab = vocab67        self.word2id = {w: i for i, w in enumerate(vocab)}68        self.dim = dim69        self.lr = lr70        n = len(vocab)71        scale = 1.0 / math.sqrt(dim)72        self.emb_in  = [[random.uniform(-scale, scale) for _ in range(dim)] for _ in range(n)]73        self.emb_out = [[0.0] * dim for _ in range(n)]7475    def _sigmoid(self, x: float) -> float:76        return 1.0 / (1.0 + math.exp(-max(-30, min(30, x))))7778    def _update(self, center_id: int, ctx_id: int, label: float) -> None:79        u = self.emb_in[center_id]80        v = self.emb_out[ctx_id]81        score = sum(a * b for a, b in zip(u, v))82        grad = (self._sigmoid(score) - label)83        for k in range(self.dim):84            self.emb_out[ctx_id][k] -= self.lr * grad * u[k]85            self.emb_in[center_id][k] -= self.lr * grad * v[k]8687    def train(self, walks: list[list[str]], window: int = 3, neg_k: int = 5) -> None:88        """Skip-Gram + 负采样训练。"""89        n = len(self.vocab)90        for walk in walks:91            ids = [self.word2id[w] for w in walk if w in self.word2id]92            for i, center in enumerate(ids):93                ctx_range = ids[max(0, i - window): i] + ids[i + 1: i + window + 1]94                for ctx in ctx_range:95                    # 正样本96                    self._update(center, ctx, 1.0)97                    # 负样本98                    for _ in range(neg_k):99                        neg = random.randint(0, n - 1)100                        while neg == ctx:101                            neg = random.randint(0, n - 1)102                        self._update(center, neg, 0.0)103104    def embedding(self, node: str) -> list[float] | None:105        idx = self.word2id.get(node)106        return self.emb_in[idx] if idx is not None else None107108109# ── 4. 召回 ──────────────────────────────────────────────────────────────────110111def cosine_sim(a: list[float], b: list[float]) -> float:112    dot = sum(x * y for x, y in zip(a, b))113    na = math.sqrt(sum(x * x for x in a))114    nb = math.sqrt(sum(x * x for x in b))115    return dot / (na * nb + 1e-9)116117118def recall_deepwalk(119    user_id: str,120    model: SkipGram,121    item_ids: list[str],122    seen: set[str],123    top_k: int = 10,124) -> list[tuple[str, float]]:125    """用用户 Embedding 检索最相似物品。"""126    u_emb = model.embedding(f"u:{user_id}")127    if u_emb is None:128        return []129    scores = []130    for item in item_ids:131        if item in seen:132            continue133        i_emb = model.embedding(f"i:{item}")134        if i_emb is not None:135            scores.append((item, cosine_sim(u_emb, i_emb)))136    scores.sort(key=lambda x: x[1], reverse=True)137    return scores[:top_k]138139140# ── 快速验证 ──────────────────────────────────────────────────────────────────141if __name__ == "__main__":142    random.seed(42)143    user_items = {144        "u1": ["i1", "i2", "i3", "i4"],145        "u2": ["i1", "i2", "i5"],146        "u3": ["i3", "i4", "i6"],147        "u4": ["i2", "i4", "i7"],148        "u5": ["i5", "i6", "i7"],149    }150    graph = build_graph(user_items)151    walks = generate_walks(graph, num_walks=20, walk_len=10)152153    vocab = list(graph.keys())154    model = SkipGram(vocab, dim=16, lr=0.05)155    model.train(walks, window=2, neg_k=3)156157    all_items = [f"i{j}" for j in range(1, 8)]158    seen = {f"i{j}" for j in [1, 2, 3, 4]}159    result = recall_deepwalk("u1", model, all_items, seen, top_k=3)160    print("DeepWalk 召回结果（u1）：")161    for item, s in result:162        print(f"  {item}: {s:.4f}")