""" AncestorTable – 固定 LUT + 可学习 ancestor embedding。 LUT 文件格式（ancestor_lut_*.pt）： { "top_k": int, "temp": float, "level_1": { "indices": Tensor [V, top_k] int64, "probs": Tensor [V, top_k] float32, }, "level_2": { ... }, # 可选，支持多层 ... } prototype 文件格式（hierarchy_prototypes_*.pt）： Tensor [K, d] 或 {"level_1": Tensor [K, d], ...} 使用方式： table = AncestorTable.from_files(lut_path, proto_path, embed_dim) table.to(device) # 加噪：采样祖先 embedding noisy = table.sample_ancestor_emb(flat_ids, level=1) # [N, d] # 全词表投影矩阵（用于 L_ancestor） W = table.projection_matrix(level=1) # [V, K]，buffer，无梯度 """ from __future__ import annotations from pathlib import Path from typing import Dict, List, Optional, Union import torch import torch.nn as nn import torch.nn.functional as F class AncestorTable(nn.Module): """ 固定祖先 LUT + 可学习 ancestor embedding。 Args: lut_indices: list of [V, top_k] int64 – 每层的祖先 index，固定 lut_probs: list of [V, top_k] float32 – 每层的祖先权重，固定 init_embeddings: list of [K_l, d] – 每层祖先 embedding 初始值 """ def __init__( self, lut_indices: List[torch.Tensor], lut_probs: List[torch.Tensor], init_embeddings: List[torch.Tensor], ): super().__init__() assert len(lut_indices) == len(lut_probs) == len(init_embeddings), \ "lut_indices, lut_probs, init_embeddings must have the same number of levels" self.num_levels = len(lut_indices) # 祖先层数（不含 leaf 和 mask） # 固定 LUT（buffer，不参与梯度） for l, (idx, prob) in enumerate(zip(lut_indices, lut_probs)): self.register_buffer(f"_lut_indices_{l}", idx.long()) self.register_buffer(f"_lut_probs_{l}", prob.float()) # 可学习 ancestor embeddings self.embeddings = nn.ParameterList([ nn.Parameter(emb.float()) for emb in init_embeddings ]) # 预计算 projection_matrix W [V, K] 并缓存（在第一次 _build_W 时填充） # 因为 W 在 lut 固定后不变，构造时一次性建好 # 注意：一律在 CPU 上构造 W，之后由 .to(device) 统一搬运，避免 CPU/GPU 索引混用。 for l in range(self.num_levels): idx = lut_indices[l].detach().cpu().long() # [V, top_k] on CPU prob = lut_probs[l].detach().cpu().float() # [V, top_k] on CPU V = idx.shape[0] K = init_embeddings[l].shape[0] W = torch.zeros(V, K, dtype=torch.float32) v_idx = torch.arange(V).unsqueeze(1).expand_as(idx) W[v_idx.reshape(-1), idx.reshape(-1)] += prob.reshape(-1) self.register_buffer(f"_W_{l}", W) # ------------------------------------------------------------------ # # Accessors # # ------------------------------------------------------------------ # def lut_indices(self, level: int) -> torch.Tensor: """返回第 level 层（1-indexed）的 [V, top_k] int64 LUT index。""" return getattr(self, f"_lut_indices_{level - 1}") def lut_probs(self, level: int) -> torch.Tensor: """返回第 level 层（1-indexed）的 [V, top_k] float32 LUT probs。""" return getattr(self, f"_lut_probs_{level - 1}") def projection_matrix(self, level: int) -> torch.Tensor: """返回第 level 层（1-indexed）的 [V, K] 全词表投影矩阵 W（buffer，无梯度）。""" return getattr(self, f"_W_{level - 1}") def ancestor_embeddings(self, level: int) -> torch.Tensor: """返回第 level 层（1-indexed）的 [K, d] 可学习 embedding。""" return self.embeddings[level - 1] # ------------------------------------------------------------------ # # Noising # # ------------------------------------------------------------------ # def sample_ancestor_emb( self, flat_ids: torch.Tensor, level: int, ) -> torch.Tensor: """ 对 flat_ids [N] 中每个 token，按 LUT probs 多项式采样一个祖先， 返回对应的可学习 embedding [N, d]。 Args: flat_ids: [N] int64 – token id level: 1-indexed 层编号 """ idx = self.lut_indices(level)[flat_ids] # [N, top_k] prob = self.lut_probs(level)[flat_ids] # [N, top_k] sampled_local = torch.multinomial(prob, num_samples=1).squeeze(1) # [N] N = flat_ids.shape[0] sampled_global = idx[torch.arange(N, device=flat_ids.device), sampled_local] # [N] emb = self.embeddings[level - 1] # [K, d] return emb[sampled_global] # [N, d] # ------------------------------------------------------------------ # # Factory # # ------------------------------------------------------------------ # @classmethod def from_files( cls, lut_path: Union[str, Path], proto_path: Optional[Union[str, Path]], embed_dim: int, device: Optional[torch.device] = None, ) -> "AncestorTable": """ 从文件构建 AncestorTable。 Args: lut_path: ancestor_lut_*.pt proto_path: hierarchy_prototypes_*.pt（None 则随机初始化） embed_dim: embedding 维度（仅 proto_path=None 时用于随机初始化） device: 加载设备 """ map_loc = device if device is not None else "cpu" # ---- 加载 LUT ---- lut = torch.load(str(lut_path), map_location=map_loc) assert isinstance(lut, dict), f"Expected dict in {lut_path}, got {type(lut)}" # 收集所有 level_N 键，按编号排序 level_keys = sorted( [k for k in lut.keys() if k.startswith("level_")], key=lambda k: int(k.split("_")[1]), ) assert len(level_keys) >= 1, f"No level_N keys found in {lut_path}" lut_indices_list = [] lut_probs_list = [] K_per_level = [] for lk in level_keys: ld = lut[lk] lut_indices_list.append(ld["indices"]) # [V, top_k] lut_probs_list.append(ld["probs"]) # [V, top_k] K_per_level.append(ld["indices"].max().item() + 1) # 推断 K # ---- 加载 / 初始化 prototype embedding ---- init_embeddings = [] if proto_path is not None and Path(str(proto_path)).exists(): proto_data = torch.load(str(proto_path), map_location=map_loc) if isinstance(proto_data, torch.Tensor): # 单层：直接是 [K, d] tensor assert len(level_keys) == 1, \ "proto_path is a single Tensor but LUT has multiple levels" init_embeddings.append(proto_data.float()) elif isinstance(proto_data, dict): for lk in level_keys: if lk in proto_data: init_embeddings.append(proto_data[lk].float()) else: # fallback：随机初始化 K = K_per_level[len(init_embeddings)] init_embeddings.append(torch.randn(K, embed_dim) * 0.02) else: raise ValueError(f"Unsupported proto_path format: {type(proto_data)}") else: # 随机初始化 for K in K_per_level: init_embeddings.append(torch.randn(K, embed_dim) * 0.02) return cls(lut_indices_list, lut_probs_list, init_embeddings)