| """ | |
| pip install pyyaml datasets wandb transformers | |
| """ | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from typing import Optional, Tuple | |
| class PositionalEncoding(nn.Module): | |
| def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1): | |
| super().__init__() | |
| self.dropout = nn.Dropout(dropout) | |
| pe = torch.zeros(max_len, d_model) | |
| position = torch.arange(0, max_len).unsqueeze(1).float() | |
| 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) | |
| pe[:, 1::2] = torch.cos(position * div_term) | |
| self.register_buffer("pe", pe.unsqueeze(0)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = x + self.pe[:, : x.size(1)] | |
| return self.dropout(x) | |
| class LoopedTransformerAutoEncoder(nn.Module): | |
| """ | |
| Encoder-decoder autoencoder where both encoder and decoder | |
| consist of a single transformer layer applied in a loop. | |
| Encoder: bidirectional, single layer looped n_encoder_loops times. | |
| Decoder: causal, single layer looped n_decoder_loops times. | |
| Memory: appended memory tokens extracted from encoder output. | |
| """ | |
| def __init__( | |
| self, | |
| vocab_size: int, | |
| d_model: int = 128, | |
| n_heads: int = 4, | |
| d_ff: int = 512, | |
| n_encoder_loops: int = 6, | |
| n_decoder_loops: int = 6, | |
| num_memory_tokens: int = 32, | |
| dropout: float = 0.0, | |
| max_seq_len: int = 512, | |
| pad_id: int = 0, | |
| bos_id: int = 1, | |
| memory_conditioning: str = "prefix", | |
| slot_source: str = "appended", | |
| ): | |
| super().__init__() | |
| self.d_model = d_model | |
| self.vocab_size = vocab_size | |
| self.num_memory_tokens = num_memory_tokens | |
| self.n_encoder_loops = n_encoder_loops | |
| self.n_decoder_loops = n_decoder_loops | |
| self.pad_id = pad_id | |
| self.bos_id = bos_id | |
| if memory_conditioning not in {"prefix", "aligned"}: | |
| raise ValueError(f"Unknown memory conditioning: {memory_conditioning}") | |
| if slot_source not in {"appended", "token_states"}: | |
| raise ValueError(f"Unknown memory slot source: {slot_source}") | |
| self.memory_conditioning = memory_conditioning | |
| self.slot_source = slot_source | |
| self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=pad_id) | |
| self.pos_encoding = PositionalEncoding(d_model, max_seq_len, dropout) | |
| self.encoder_layer = nn.TransformerEncoderLayer( | |
| d_model, n_heads, d_ff, dropout, batch_first=True, norm_first=True | |
| ) | |
| self.encoder_norm = nn.LayerNorm(d_model) | |
| self.decoder_layer = nn.TransformerEncoderLayer( | |
| d_model, n_heads, d_ff, dropout, batch_first=True, norm_first=True | |
| ) | |
| self.decoder_norm = nn.LayerNorm(d_model) | |
| self.enc_proj = nn.Linear(d_model, d_model) | |
| self.enc_to_token = nn.Linear(d_model, d_model) | |
| self.dropout_layer = nn.Dropout(dropout) | |
| self.lm_head = nn.Linear(d_model, vocab_size, bias=False) | |
| self._init_weights() | |
| def _init_weights(self): | |
| for module in [self.enc_proj, self.enc_to_token]: | |
| nn.init.xavier_uniform_(module.weight) | |
| nn.init.zeros_(module.bias) | |
| nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.02) | |
| def encode( | |
| self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None | |
| ) -> torch.Tensor: | |
| x = self.pos_encoding(self.embedding(input_ids) * math.sqrt(self.d_model)) | |
| src_key_padding_mask = None | |
| if attention_mask is not None: | |
| src_key_padding_mask = (attention_mask == 0) | |
| for _ in range(self.n_encoder_loops): | |
| x = self.encoder_layer(x, src_key_padding_mask=src_key_padding_mask) | |
| x = self.encoder_norm(x) | |
| if self.slot_source == "appended": | |
| memory_hidden = x[:, -self.num_memory_tokens :, :] | |
| else: | |
| if x.size(1) != self.num_memory_tokens: | |
| raise ValueError( | |
| "token_states requires one memory slot per source position: " | |
| f"source length={x.size(1)}, memory slots={self.num_memory_tokens}" | |
| ) | |
| memory_hidden = x | |
| return self.enc_proj(memory_hidden) | |
| def decode_training( | |
| self, | |
| enc_vec: torch.Tensor, | |
| target_ids: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| ) -> torch.Tensor: | |
| B, L = target_ids.shape | |
| memory = self.enc_to_token(enc_vec) | |
| K = memory.size(1) | |
| bos_ids = torch.full( | |
| (B, 1), self.bos_id, device=target_ids.device, dtype=target_ids.dtype | |
| ) | |
| teacher_ids = torch.cat([bos_ids, target_ids[:, :-1]], dim=1) | |
| teacher_embeds = self.embedding(teacher_ids) * math.sqrt(self.d_model) | |
| if self.memory_conditioning == "aligned": | |
| if K != L: | |
| raise ValueError(f"aligned conditioning requires K == L, got K={K}, L={L}") | |
| combined = self.pos_encoding(memory + teacher_embeds) | |
| output_start = 0 | |
| else: | |
| # Decoder positions cover the complete [memory prefix, shifted targets] | |
| # sequence. Applying positions before concatenation incorrectly gives | |
| # the first target the same position as the first memory slot. | |
| combined = self.pos_encoding(torch.cat([memory, teacher_embeds], dim=1)) | |
| output_start = K | |
| total_len = combined.size(1) | |
| causal_mask = torch.triu( | |
| torch.ones(total_len, total_len, device=combined.device), diagonal=1 | |
| ).bool() | |
| padding_mask = None | |
| if attention_mask is not None: | |
| if self.memory_conditioning == "aligned": | |
| padding_mask = attention_mask == 0 | |
| else: | |
| prefix_mask = torch.zeros(B, K, device=attention_mask.device, dtype=attention_mask.dtype) | |
| combined_pad = torch.cat([prefix_mask, (attention_mask == 0).long()], dim=1) | |
| padding_mask = combined_pad.bool() | |
| x = combined | |
| for _ in range(self.n_decoder_loops): | |
| x = self.decoder_layer( | |
| x, src_mask=causal_mask, src_key_padding_mask=padding_mask | |
| ) | |
| x = self.decoder_norm(x) | |
| hidden = self.dropout_layer(x[:, output_start : output_start + L]) | |
| return self.lm_head(hidden) | |
| def forward( | |
| self, | |
| input_ids: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| target_ids: Optional[torch.Tensor] = None, | |
| target_attention_mask: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| enc_vec = self.encode(input_ids, attention_mask) | |
| target_ids = input_ids if target_ids is None else target_ids | |
| target_attention_mask = attention_mask if target_attention_mask is None else target_attention_mask | |
| logits = self.decode_training(enc_vec, target_ids, target_attention_mask) | |
| return logits, enc_vec | |
| def generate( | |
| self, | |
| input_ids: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| max_length: int = 100, | |
| temperature: float = 1.0, | |
| do_sample: bool = False, | |
| ) -> torch.Tensor: | |
| self.eval() | |
| enc_vec = self.encode(input_ids, attention_mask) | |
| memory = self.enc_to_token(enc_vec) | |
| B = input_ids.size(0) | |
| generated_ids = torch.full( | |
| (B, 1), self.bos_id, device=input_ids.device, dtype=input_ids.dtype | |
| ) | |
| for _ in range(max_length): | |
| token_embeds = self.embedding(generated_ids) * math.sqrt(self.d_model) | |
| if self.memory_conditioning == "aligned": | |
| steps = token_embeds.size(1) | |
| if steps > memory.size(1): | |
| raise ValueError("generation length exceeds aligned memory slots") | |
| decoder_input = self.pos_encoding(token_embeds + memory[:, :steps]) | |
| else: | |
| decoder_input = self.pos_encoding(torch.cat([memory, token_embeds], dim=1)) | |
| total_len = decoder_input.size(1) | |
| causal_mask = torch.triu( | |
| torch.ones(total_len, total_len, device=input_ids.device), diagonal=1 | |
| ).bool() | |
| x = decoder_input | |
| for _ in range(self.n_decoder_loops): | |
| x = self.decoder_layer(x, src_mask=causal_mask) | |
| x = self.decoder_norm(x) | |
| logits = self.lm_head(x[:, -1]) / temperature | |
| if do_sample: | |
| next_id = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) | |
| else: | |
| next_id = logits.argmax(dim=-1, keepdim=True) | |
| generated_ids = torch.cat([generated_ids, next_id], dim=1) | |
| return generated_ids[:, 1:] | |
Xet Storage Details
- Size:
- 8.99 kB
- Xet hash:
- b2c7fb1a91e433835bd12dd36831b7a92b2c68a11ec016c78884c8c02743f152
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.