| from __future__ import annotations |
|
|
| from typing import Dict, List, Optional, Tuple |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from ..tokenizer.bpe_standard import BPETokenizer |
|
|
|
|
| class TokenizerVocabAdapter: |
| """Bridges tokenizer vocabulary to integer token IDs.""" |
|
|
| def __init__( |
| self, |
| tokenizer: BPETokenizer, |
| pad_token: str = "<pad>", |
| unk_token: str = "<unk>", |
| bos_token: Optional[str] = None, |
| eos_token: Optional[str] = None, |
| ): |
| self.tokenizer = tokenizer |
| self.vocab = tokenizer.vocab |
| self.reverse_vocab = {v: k for k, v in self.vocab.items()} |
|
|
| self.pad_token = pad_token |
| self.unk_token = unk_token |
| self.bos_token = bos_token |
| self.eos_token = eos_token |
|
|
| |
| self._ensure_special_tokens() |
|
|
| def _ensure_special_tokens(self) -> None: |
| """Register special tokens in vocab if missing.""" |
| specials = [ |
| (self.pad_token, 0), |
| (self.unk_token, 1), |
| ] |
| if self.bos_token: |
| specials.append((self.bos_token, 2)) |
| if self.eos_token: |
| specials.append((self.eos_token, 3 if self.bos_token else 2)) |
|
|
| for token, fallback_id in specials: |
| if token not in self.vocab: |
| new_id = max(self.vocab.values()) + 1 if self.vocab else fallback_id |
| self.vocab[token] = new_id |
| self.reverse_vocab[new_id] = token |
|
|
| def get_vocab_size(self) -> int: |
| return len(self.vocab) |
|
|
| def token_to_id(self, token: str) -> int: |
| return self.vocab.get(token, self.vocab.get(self.unk_token, 1)) |
|
|
| def id_to_token(self, token_id: int) -> str: |
| return self.reverse_vocab.get(token_id, self.unk_token) |
|
|
| def tokens_to_ids(self, tokens: List[str]) -> List[int]: |
| return [self.token_to_id(token) for token in tokens] |
|
|
| def ids_to_tokens(self, token_ids: List[int]) -> List[str]: |
| return [self.id_to_token(token_id) for token_id in token_ids] |
|
|
| def get_special_token_ids(self) -> Dict[str, int]: |
| return { |
| "pad": self.vocab.get(self.pad_token, 0), |
| "unk": self.vocab.get(self.unk_token, 1), |
| "bos": self.vocab.get(self.bos_token, None) if self.bos_token else None, |
| "eos": self.vocab.get(self.eos_token, None) if self.eos_token else None, |
| } |
|
|
|
|
| class PositionalEmbedding(nn.Module): |
| """Learned positional embeddings.""" |
|
|
| def __init__(self, d_model: int, max_seq_length: int = 2048): |
| super().__init__() |
| self.d_model = d_model |
| self.max_seq_length = max_seq_length |
| self.pos_embed = nn.Embedding(max_seq_length, d_model) |
|
|
| def forward(self, seq_len: int) -> torch.Tensor: |
| """Return positional embeddings for sequence length. |
| |
| Args: |
| seq_len: sequence length |
| |
| Returns: |
| pos_embed: shape [seq_len, d_model] |
| """ |
| positions = torch.arange(seq_len, dtype=torch.long, device=self.pos_embed.weight.device) |
| return self.pos_embed(positions) |
|
|
|
|
| class TokenEmbedding(nn.Module): |
| """Simple token embedding without factorization.""" |
|
|
| 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.embed = nn.Embedding(vocab_size, d_model, padding_idx=padding_idx) |
| self.scale = (d_model ** 0.5) |
|
|
| def forward(self, token_ids: torch.Tensor) -> torch.Tensor: |
| return self.embed(token_ids) * self.scale |
|
|
