| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from yaml_bert.config import TreePosVariant, YamlBertConfig |
|
|
|
|
| class YamlBertEmbedding(nn.Module): |
| """v9 embedding layer: single subword table + tree positional encoding. |
| |
| Produces input vectors by summing: |
| - Subword embedding (looked up by token_id; same table for KEY and VALUE |
| positions — what they ARE is signalled separately via node_type_emb) |
| - Tree positional encoding (composition depends on config.tree_pos_variant) |
| """ |
|
|
| def __init__( |
| self, |
| config: YamlBertConfig, |
| subword_vocab_size: int, |
| ) -> None: |
| super().__init__() |
| d: int = config.d_model |
| variant: TreePosVariant = config.tree_pos_variant |
| self.variant: TreePosVariant = variant |
|
|
| self.subword_embedding: nn.Embedding = nn.Embedding(subword_vocab_size, d) |
| self.node_type_embedding: nn.Embedding = nn.Embedding(4, d) |
|
|
| use_depth: bool = variant in (TreePosVariant.FULL, TreePosVariant.NO_SIBLING) |
| use_sibling: bool = variant in (TreePosVariant.FULL, TreePosVariant.NO_DEPTH) |
| use_seq_pos: bool = variant == TreePosVariant.SEQUENTIAL |
|
|
| self.depth_embedding: nn.Embedding | None = ( |
| nn.Embedding(config.max_depth, d) if use_depth else None |
| ) |
| self.sibling_embedding: nn.Embedding | None = ( |
| nn.Embedding(config.max_sibling, d) if use_sibling else None |
| ) |
| self.pos_embedding: nn.Embedding | None = ( |
| nn.Embedding(config.max_seq_len, d) if use_seq_pos else None |
| ) |
|
|
| self.layer_norm: nn.LayerNorm = nn.LayerNorm(d) |
|
|
| def forward( |
| self, |
| token_ids: torch.Tensor, |
| node_types: torch.Tensor, |
| depths: torch.Tensor, |
| sibling_indices: torch.Tensor, |
| ) -> torch.Tensor: |
| token_emb = self.subword_embedding(token_ids) |
|
|
| tree_pos = self.node_type_embedding(node_types) |
| if self.depth_embedding is not None: |
| tree_pos = tree_pos + self.depth_embedding(depths) |
| if self.sibling_embedding is not None: |
| tree_pos = tree_pos + self.sibling_embedding(sibling_indices) |
| if self.pos_embedding is not None: |
| seq_len: int = token_ids.size(1) |
| max_pos: int = self.pos_embedding.num_embeddings |
| positions = ( |
| torch.arange(seq_len, device=token_ids.device) |
| .clamp(max=max_pos - 1) |
| .unsqueeze(0) |
| .expand(token_ids.size(0), seq_len) |
| ) |
| tree_pos = tree_pos + self.pos_embedding(positions) |
|
|
| return self.layer_norm(token_emb + tree_pos) |
|
|
