| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import math | |
| from typing import Optional, Tuple, List | |
| class PositionalEncoding(nn.Module): | |
| def __init__(self, d_model: int, max_seq_len: int = 5000, dropout: float = 0.1): | |
| super(PositionalEncoding, self).__init__() | |
| self.d_model = d_model | |
| self.dropout = nn.Dropout(dropout) | |
| pe = torch.zeros(max_seq_len, d_model) | |
| position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1) | |
| div_term = torch.exp(torch.arange(0, d_model, 2).float() * | |
| -(math.log(10000.0) / d_model)) | |
| pe[:, 0::2] = torch.sin(position * div_term) | |
| if d_model % 2 == 1: | |
| pe[:, 1::2] = torch.cos(position * div_term[:-1]) | |
| else: | |
| pe[:, 1::2] = torch.cos(position * div_term) | |
| self.register_buffer('pe', pe.unsqueeze(0)) | |
| def forward(self, x): | |
| batch_size, seq_len, d_model = x.size() | |
| x = x + self.pe[:, :seq_len, :d_model] | |
| return self.dropout(x) | |
| class LearnedPositionalEmbedding(nn.Module): | |
| def __init__(self, max_seq_len: int, d_model: int, dropout: float = 0.1): | |
| super(LearnedPositionalEmbedding, self).__init__() | |
| self.max_seq_len = max_seq_len | |
| self.d_model = d_model | |
| self.pos_embedding = nn.Embedding(max_seq_len, d_model) | |
| self.dropout = nn.Dropout(dropout) | |
| nn.init.normal_(self.pos_embedding.weight, std=0.02) | |
| def forward(self, x): | |
| batch_size, seq_len, d_model = x.size() | |
| if seq_len > self.max_seq_len: | |
| raise ValueError(f"Sequence length {seq_len} exceeds maximum {self.max_seq_len}") | |
| positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1) | |
| pos_emb = self.pos_embedding(positions) | |
| x = x + pos_emb | |
| return self.dropout(x) | |
| class RotaryPositionalEmbedding(nn.Module): | |
| def __init__(self, d_model: int, max_seq_len: int = 2048, base: float = 10000.0): | |
| super(RotaryPositionalEmbedding, self).__init__() | |
| self.d_model = d_model | |
| self.max_seq_len = max_seq_len | |
| self.base = base | |
| inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model)) | |
| self.register_buffer('inv_freq', inv_freq) | |
| self._seq_len_cached = 0 | |
| self._cos_cached = None | |
| self._sin_cached = None | |
| def _update_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype): | |
| if seq_len > self._seq_len_cached: | |
| self._seq_len_cached = seq_len | |
| t = torch.arange(seq_len, device=device, dtype=torch.float32) | |
| freqs = torch.outer(t, self.inv_freq) | |
| self._cos_cached = freqs.cos().to(dtype) | |
| self._sin_cached = freqs.sin().to(dtype) | |
| def forward(self, q: torch.Tensor, k: torch.Tensor, start_pos: int = 0) -> Tuple[torch.Tensor, torch.Tensor]: | |
| batch_size, seq_len, num_heads, head_dim = q.shape | |
| self._update_cos_sin_cache(start_pos + seq_len, q.device, q.dtype) | |
| cos = self._cos_cached[start_pos:start_pos + seq_len, :head_dim // 2] | |
| sin = self._sin_cached[start_pos:start_pos + seq_len, :head_dim // 2] | |
| cos = cos.view(1, seq_len, 1, -1) | |
| sin = sin.view(1, seq_len, 1, -1) | |
| q = q.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim) | |
| k = k.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim) | |
| q_rot = self._rotate_half(q, cos, sin) | |
| k_rot = self._rotate_half(k, cos, sin) | |
| q_rot = q_rot.reshape(batch_size, num_heads, seq_len, head_dim).transpose(1, 2) | |
| k_rot = k_rot.reshape(batch_size, num_heads, seq_len, head_dim).transpose(1, 2) | |
| return q_rot, k_rot | |
| def _rotate_half(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: | |
| x1 = x[..., :x.shape[-1] // 2] | |
| x2 = x[..., x.shape[-1] // 2:] | |
| return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) | |
| class TechEmbeddingLayer(nn.Module): | |
| def __init__(self, | |
| vocab_size: int, | |
| d_model: int, | |
| max_seq_len: int = 512, | |
| dropout: float = 0.1, | |
| padding_idx: int = 0, | |
| pos_encoding: str = "learned", | |
| layer_norm: bool = True): | |
| super(TechEmbeddingLayer, self).__init__() | |
| self.d_model = d_model | |
| self.vocab_size = vocab_size | |
| self.padding_idx = padding_idx | |
| self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_idx) | |
| self.pos_encoding_type = pos_encoding | |
| if pos_encoding == "sinusoidal": | |
| self.pos_encoding = PositionalEncoding(d_model, max_seq_len, dropout) | |
| elif pos_encoding == "learned": | |
| self.pos_encoding = LearnedPositionalEmbedding(max_seq_len, d_model, dropout) | |
| elif pos_encoding == "rope": | |
| self.pos_encoding = RotaryPositionalEmbedding(d_model, max_seq_len) | |
| else: | |
| raise ValueError(f"Unknown positional encoding type: {pos_encoding}") | |
| self.layer_norm = nn.LayerNorm(d_model) if layer_norm else nn.Identity() | |
| self.dropout = nn.Dropout(dropout) | |
| self._init_weights() | |
| def _init_weights(self): | |
| nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02) | |
| if self.padding_idx is not None: | |
| nn.init.constant_(self.token_embedding.weight[self.padding_idx], 0.0) | |
| def forward(self, input_ids: torch.Tensor) -> torch.Tensor: | |
| if (input_ids >= self.vocab_size).any(): | |
| raise ValueError(f"Input IDs contain values >= vocab_size ({self.vocab_size})") | |
| embeddings = self.token_embedding(input_ids) | |
| if self.pos_encoding_type != "rope": | |
| embeddings = self.pos_encoding(embeddings) | |
| embeddings = self.layer_norm(embeddings) | |
| embeddings = self.dropout(embeddings) | |
| return embeddings | |
| def get_positional_encoding(self): | |
| return self.pos_encoding if self.pos_encoding_type == "rope" else None | |
| class AdaptiveEmbedding(nn.Module): | |
| def __init__(self, | |
| vocab_size: int, | |
| d_model: int, | |
| cutoffs: list = [2000, 10000], | |
| div_val: float = 4.0): | |
| super(AdaptiveEmbedding, self).__init__() | |
| self.vocab_size = vocab_size | |
| self.d_model = d_model | |
| self.cutoffs = [0] + cutoffs + [vocab_size] | |
| self.div_val = div_val | |
| self.embeddings = nn.ModuleList() | |
| self.projections = nn.ModuleList() | |
| for i in range(len(self.cutoffs) - 1): | |
| l_idx = self.cutoffs[i] | |
| r_idx = self.cutoffs[i + 1] | |
| d_emb = int(d_model / (div_val ** i)) | |
| emb = nn.Embedding(r_idx - l_idx, d_emb) | |
| nn.init.normal_(emb.weight, mean=0.0, std=0.02) | |
| self.embeddings.append(emb) | |
| if d_emb != d_model: | |
| proj = nn.Linear(d_emb, d_model, bias=False) | |
| nn.init.normal_(proj.weight, mean=0.0, std=0.02) | |
| self.projections.append(proj) | |
| else: | |
| self.projections.append(nn.Identity()) | |
| def forward(self, input_ids: torch.Tensor) -> torch.Tensor: | |
| if (input_ids >= self.vocab_size).any(): | |
| raise ValueError(f"Input IDs contain values >= vocab_size ({self.vocab_size})") | |
| batch_size, seq_len = input_ids.shape | |
| embeddings = torch.zeros(batch_size, seq_len, self.d_model, | |
| device=input_ids.device, dtype=torch.float32) | |
| for i in range(len(self.cutoffs) - 1): | |
| l_idx = self.cutoffs[i] | |
| r_idx = self.cutoffs[i + 1] | |
| mask = (input_ids >= l_idx) & (input_ids < r_idx) | |
| if mask.any(): | |
| indices = input_ids[mask] - l_idx | |
| indices = indices.clamp(max=r_idx - l_idx - 1) | |
| emb = self.embeddings[i](indices) | |
| emb = self.projections[i](emb) | |
| embeddings[mask] = emb | |
| return embeddings | |
| def create_padding_mask(input_ids: torch.Tensor, padding_idx: int = 0) -> torch.Tensor: | |
| return input_ids == padding_idx | |
| def create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor: | |
| return torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool() | |
| def create_attention_mask(input_ids: torch.Tensor, | |
| padding_idx: int = 0, | |
| causal: bool = True) -> torch.Tensor: | |
| batch_size, seq_len = input_ids.shape | |
| device = input_ids.device | |
| padding_mask = create_padding_mask(input_ids, padding_idx) | |
| padding_mask = padding_mask.unsqueeze(1).expand(batch_size, seq_len, seq_len) | |
| if causal: | |
| causal_mask = create_causal_mask(seq_len, device) | |
| causal_mask = causal_mask.unsqueeze(0).expand(batch_size, seq_len, seq_len) | |
| combined_mask = padding_mask | causal_mask | |
| else: | |
| combined_mask = padding_mask | |
| return combined_mask | |
| class EmbeddingAnalyzer: | |
| def __init__(self, embedding_layer: nn.Module): | |
| self.embedding_layer = embedding_layer | |
| def get_similarity_matrix(self, tokens: List[int] = None) -> torch.Tensor: | |
| if hasattr(self.embedding_layer, 'token_embedding'): | |
| embeddings = self.embedding_layer.token_embedding.weight | |
| elif hasattr(self.embedding_layer, 'embeddings'): | |
| weights = [emb.weight for emb in self.embedding_layer.embeddings] | |
| embeddings = [] | |
| for i, w in enumerate(weights): | |
| proj = self.embedding_layer.projections[i] | |
| embeddings.append(proj(w)) | |
| embeddings = torch.cat(embeddings, dim=0) | |
| else: | |
| embeddings = self.embedding_layer.weight | |
| if tokens is not None and len(tokens) > 0: | |
| embeddings = embeddings[tokens] | |
| normalized_embeddings = F.normalize(embeddings, p=2, dim=1) | |
| return torch.mm(normalized_embeddings, normalized_embeddings.t()) | |
| def find_similar_tokens(self, token_id: int, top_k: int = 10) -> List[Tuple[int, float]]: | |
| similarity_matrix = self.get_similarity_matrix() | |
| similarities = similarity_matrix[token_id] | |
| top_similarities, top_indices = torch.topk(similarities, top_k + 1) | |
| mask = top_indices != token_id | |
| top_similarities = top_similarities[mask][:top_k] | |
| top_indices = top_indices[mask][:top_k] | |
| return list(zip(top_indices.tolist(), top_similarities.tolist())) | |
| def analyze_embedding_distribution(self): | |
| if hasattr(self.embedding_layer, 'token_embedding'): | |
| weights = self.embedding_layer.token_embedding.weight | |
| elif hasattr(self.embedding_layer, 'embeddings'): | |
| weights = torch.cat([emb.weight for emb in self.embedding_layer.embeddings], dim=0) | |
| else: | |
| weights = self.embedding_layer.weight | |
| stats = { | |
| 'mean': weights.mean().item(), | |
| 'std': weights.std().item(), | |
| 'min': weights.min().item(), | |
| 'max': weights.max().item(), | |
| 'norm_mean': weights.norm(dim=1).mean().item(), | |
| 'norm_std': weights.norm(dim=1).std().item() | |
| } | |
| return stats | |
| def test_embeddings(): | |
| print("Testing embedding layers...") | |
| vocab_size = 1000 | |
| d_model = 512 | |
| max_seq_len = 128 | |
| batch_size = 4 | |
| seq_len = 64 | |
| input_ids = torch.randint(1, vocab_size, (batch_size, seq_len)) | |
| embedding_types = [ | |
| ("Learned Position", "learned"), | |
| ("Sinusoidal Position", "sinusoidal"), | |
| ("RoPE", "rope") | |
| ] | |
| for name, pos_type in embedding_types: | |
| print(f"\nTesting {name} Embedding:") | |
| embedding_layer = TechEmbeddingLayer( | |
| vocab_size=vocab_size, | |
| d_model=d_model, | |
| max_seq_len=max_seq_len, | |
| pos_encoding=pos_type | |
| ) | |
| embeddings = embedding_layer(input_ids) | |
| print(f"Input shape: {input_ids.shape}") | |
| print(f"Output shape: {embeddings.shape}") | |
| print(f"Expected shape: ({batch_size}, {seq_len}, {d_model})") | |
| analyzer = EmbeddingAnalyzer(embedding_layer) | |
| stats = analyzer.analyze_embedding_distribution() | |
| print(f"Embedding statistics:") | |
| for key, value in stats.items(): | |
| print(f" {key}: {value:.4f}") | |
| print(f"\nTesting Adaptive Embeddings:") | |
| adaptive_emb = AdaptiveEmbedding( | |
| vocab_size=vocab_size, | |
| d_model=d_model, | |
| cutoffs=[200, 500], | |
| div_val=2.0 | |
| ) | |
| embeddings = adaptive_emb(input_ids) | |
| print(f"Adaptive embedding output shape: {embeddings.shape}") | |
| print(f"\nTesting masking functions:") | |
| input_ids_padded = input_ids.clone() | |
| input_ids_padded[:, -10:] = 0 | |
| padding_mask = create_padding_mask(input_ids_padded, padding_idx=0) | |
| causal_mask = create_causal_mask(seq_len, input_ids.device) | |
| attention_mask = create_attention_mask(input_ids_padded, padding_idx=0, causal=True) | |
| print(f"Padding mask shape: {padding_mask.shape}") | |
| print(f"Causal mask shape: {causal_mask.shape}") | |
| print(f"Attention mask shape: {attention_mask.shape}") | |
| print(f"Padding positions: {padding_mask.sum().item()}") | |
| print(f"Causal mask positions: {causal_mask.sum().item()}") | |
| print(f"Combined mask positions: {attention_mask.sum().item()}") | |
| print("\nAll embedding tests completed successfully!") | |
| if __name__ == "__main__": | |
| test_embeddings() |