File size: 2,414 Bytes
8594de8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """
Xavante - embedding_reconfig.py
Responsabilidade: Embedding reconfiguravel (Teorema 12.1/12.2).
Permite expansao dinamica do vocabulario e reconfiguracao de dim do modelo.
"""
from __future__ import annotations
import logging
import math
from typing import Optional
import torch
import torch.nn as nn
logger = logging.getLogger(__name__)
class ReconfigurableEmbedding(nn.Module):
"""
Embedding que suporta:
- Expansao de vocabulario sem reset
- Reconfiguracao de dim (projecao)
- Quantizacao para memoria
"""
def __init__(self, vocab_size: int, d_model: int, padding_idx: Optional[int] = None):
super().__init__()
self.vocab_size = vocab_size
self.d_model = d_model
self.padding_idx = padding_idx
self.weight = nn.Parameter(torch.empty(vocab_size, d_model))
nn.init.normal_(self.weight, mean=0.0, std=1.0 / math.sqrt(d_model) if False else 0.02)
if padding_idx is not None:
with torch.no_grad():
self.weight[padding_idx].fill_(0)
def forward(self, idx: torch.Tensor) -> torch.Tensor:
return torch.nn.functional.embedding(idx, self.weight, padding_idx=self.padding_idx)
def expand_vocab(self, new_size: int) -> None:
"""Expande o vocabulario preservando pesos antigos."""
if new_size <= self.vocab_size:
return
old = self.weight.data
new = torch.empty(new_size - self.vocab_size, self.d_model, device=old.device, dtype=old.dtype)
nn.init.normal_(new, mean=0.0, std=0.02)
self.weight = nn.Parameter(torch.cat([old, new], dim=0))
self.vocab_size = new_size
logger.info("Vocab expandido para %d", new_size)
def reconfigure_dim(self, new_d: int) -> None:
"""Reconfigura a dimensao via projecao linear (random init para a parte nova)."""
if new_d == self.d_model:
return
old = self.weight.data # [V, d_model]
if new_d > self.d_model:
pad = torch.empty(old.shape[0], new_d - self.d_model, device=old.device, dtype=old.dtype)
nn.init.normal_(pad, mean=0.0, std=0.02)
new_w = torch.cat([old, pad], dim=1)
else:
new_w = old[:, :new_d]
self.weight = nn.Parameter(new_w)
self.d_model = new_d
logger.info("Dim reconfigurada para %d", new_d)
__all__ = ["ReconfigurableEmbedding"]
|