| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import PretrainedConfig, PreTrainedModel |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
| from typing import Optional, Tuple |
| import math |
| class RecursiveLanguageModelConfig(PretrainedConfig): |
| model_type = "recursive_language_model" |
| def __init__( |
| self, |
| vocab_size: int = 50257, |
| embedding_dim: int = 512, |
| num_layers: int = 6, |
| num_attention_heads: int = 8, |
| max_recursion_steps: int = 5, |
| max_position_embeddings: int = 512, |
| hidden_dropout_prob: float = 0.1, |
| attention_dropout_prob: float = 0.1, |
| intermediate_size: int = 2048, |
| layer_norm_eps: float = 1e-5, |
| pad_token_id: int = 50256, |
| bos_token_id: int = 50256, |
| eos_token_id: int = 50256, |
| simple_recursion_steps: int = 1, |
| medium_recursion_steps: int = 3, |
| complex_recursion_steps: int = 5, |
| confidence_threshold: float = 0.8, |
| use_adaptive_stopping: bool = True, |
| initializer_range: float = 0.02, |
| **kwargs |
| ): |
| super().__init__( |
| pad_token_id=pad_token_id, |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| **kwargs |
| ) |
| self.vocab_size = vocab_size |
| self.embedding_dim = embedding_dim |
| self.num_layers = num_layers |
| self.num_attention_heads = num_attention_heads |
| self.max_recursion_steps = max_recursion_steps |
| self.max_position_embeddings = max_position_embeddings |
| self.hidden_dropout_prob = hidden_dropout_prob |
| self.attention_dropout_prob = attention_dropout_prob |
| self.intermediate_size = intermediate_size |
| self.layer_norm_eps = layer_norm_eps |
| self.simple_recursion_steps = simple_recursion_steps |
| self.medium_recursion_steps = medium_recursion_steps |
| self.complex_recursion_steps = complex_recursion_steps |
| self.confidence_threshold = confidence_threshold |
| self.use_adaptive_stopping = use_adaptive_stopping |
| self.initializer_range = initializer_range |
| class RotaryPositionalEmbedding(nn.Module): |
| def __init__(self, dim, max_seq_len=2048, base=10000): |
| super().__init__() |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer('inv_freq', inv_freq) |
| self.max_seq_len = max_seq_len |
| self.dim = dim |
| def forward(self, seq_len, device): |
| t = torch.arange(seq_len, device=device).type_as(self.inv_freq) |
| freqs = torch.einsum('i,j->ij', t, self.inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| return emb.cos(), emb.sin() |
| def apply_rotary_pos_emb(q, k, cos, sin): |
| def rotate_half(x): |
| x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:] |
| return torch.cat((-x2, x1), dim=-1) |
| q_embed = (q * cos) + (rotate_half(q) * sin) |
| k_embed = (k * cos) + (rotate_half(k) * sin) |
| return q_embed, k_embed |
| class MultiHeadAttention(nn.Module): |
| def __init__(self, config: RecursiveLanguageModelConfig): |
| super().__init__() |
| self.num_heads = config.num_attention_heads |
| self.head_dim = config.embedding_dim // config.num_attention_heads |
| self.embed_dim = config.embedding_dim |
| assert self.embed_dim % self.num_heads == 0 |
| self.q_proj = nn.Linear(config.embedding_dim, config.embedding_dim) |
| self.k_proj = nn.Linear(config.embedding_dim, config.embedding_dim) |
| self.v_proj = nn.Linear(config.embedding_dim, config.embedding_dim) |
| self.out_proj = nn.Linear(config.embedding_dim, config.embedding_dim) |
| self.dropout = nn.Dropout(config.attention_dropout_prob) |
| self.rotary_emb = RotaryPositionalEmbedding(self.head_dim, config.max_position_embeddings) |
| def forward(self, hidden_states, attention_mask=None): |
| batch_size, seq_len, _ = hidden_states.shape |
| q = self.q_proj(hidden_states) |
| k = self.k_proj(hidden_states) |
| v = self.v_proj(hidden_states) |
| q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) |
| k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) |
| v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) |
| cos, sin = self.rotary_emb(seq_len, hidden_states.device) |
| cos = cos[None, None, :, :].expand(batch_size, self.num_heads, -1, -1) |
| sin = sin[None, None, :, :].expand(batch_size, self.num_heads, -1, -1) |
| q, k = apply_rotary_pos_emb(q, k, cos, sin) |
| attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| if attention_mask is not None: |
| attn_weights = attn_weights + attention_mask |
| attn_weights = F.softmax(attn_weights, dim=-1) |
| attn_weights = self.dropout(attn_weights) |
| attn_output = torch.matmul(attn_weights, v) |
| attn_output = attn_output.transpose(1, 2).contiguous() |
| attn_output = attn_output.view(batch_size, seq_len, self.embed_dim) |
| attn_output = self.out_proj(attn_output) |
| return attn_output |
| class FeedForward(nn.Module): |
| def __init__(self, config: RecursiveLanguageModelConfig): |
| super().__init__() |
| self.fc1 = nn.Linear(config.embedding_dim, config.intermediate_size) |
| self.fc2 = nn.Linear(config.intermediate_size, config.embedding_dim) |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| def forward(self, x): |
| x = self.fc1(x) |
| x = F.gelu(x) |
| x = self.dropout(x) |
| x = self.fc2(x) |
| x = self.dropout(x) |
| return x |
| class TransformerBlock(nn.Module): |
| def __init__(self, config: RecursiveLanguageModelConfig): |
| super().__init__() |
| self.attention = MultiHeadAttention(config) |
| self.feed_forward = FeedForward(config) |
| self.ln1 = nn.LayerNorm(config.embedding_dim, eps=config.layer_norm_eps) |
| self.ln2 = nn.LayerNorm(config.embedding_dim, eps=config.layer_norm_eps) |
| def forward(self, hidden_states, attention_mask=None): |
| residual = hidden_states |
| hidden_states = self.ln1(hidden_states) |
| hidden_states = self.attention(hidden_states, attention_mask) |
| hidden_states = residual + hidden_states |
| residual = hidden_states |
| hidden_states = self.ln2(hidden_states) |
| hidden_states = self.feed_forward(hidden_states) |
| hidden_states = residual + hidden_states |
| return hidden_states |
| class SequenceLevelRouter(nn.Module): |
| def __init__(self, config: RecursiveLanguageModelConfig): |
| super().__init__() |
| self.config = config |
| self.pooler = nn.Linear(config.embedding_dim, config.embedding_dim) |
| self.pooler_activation = nn.Tanh() |
| self.classifier = nn.Sequential( |
| nn.Linear(config.embedding_dim, config.embedding_dim // 2), |
| nn.GELU(), |
| nn.Dropout(0.1), |
| nn.Linear(config.embedding_dim // 2, 3) |
| ) |
| def forward(self, hidden_states, attention_mask=None): |
| if attention_mask is not None: |
| mask_expanded = attention_mask.unsqueeze(-1).float() |
| sum_hidden = torch.sum(hidden_states * mask_expanded, dim=1) |
| sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9) |
| pooled = sum_hidden / sum_mask |
| else: |
| pooled = hidden_states.mean(dim=1) |
| pooled = self.pooler(pooled) |
| pooled = self.pooler_activation(pooled) |
| complexity_logits = self.classifier(pooled) |
| complexity_class = torch.argmax(complexity_logits, dim=-1) |
| recursion_steps = torch.zeros_like(complexity_class) |
| recursion_steps[complexity_class == 0] = self.config.simple_recursion_steps |
| recursion_steps[complexity_class == 1] = self.config.medium_recursion_steps |
| recursion_steps[complexity_class == 2] = self.config.complex_recursion_steps |
| return complexity_logits, complexity_class, recursion_steps |
| class RecursionLayer(nn.Module): |
| def __init__(self, config: RecursiveLanguageModelConfig): |
| super().__init__() |
| self.transformer_block = TransformerBlock(config) |
| def forward(self, hidden_states, attention_mask=None): |
| return self.transformer_block(hidden_states, attention_mask) |
| class RecursiveLanguageModel(PreTrainedModel): |
| config_class = RecursiveLanguageModelConfig |
| def __init__(self, config: RecursiveLanguageModelConfig): |
| super().__init__(config) |
| self.config = config |
| self.embedding_layer = nn.Embedding( |
| config.vocab_size, |
| config.embedding_dim, |
| padding_idx=config.pad_token_id |
| ) |
| self.base_transformer = nn.ModuleList([ |
| TransformerBlock(config) for _ in range(config.num_layers) |
| ]) |
| self.router = SequenceLevelRouter(config) |
| self.recursion_layer = RecursionLayer(config) |
| self.final_layer_norm = nn.LayerNorm(config.embedding_dim, eps=config.layer_norm_eps) |
| self.language_model_head = nn.Linear(config.embedding_dim, config.vocab_size, bias=False) |
| self.post_init() |
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) |
| if module.bias is not None: |
| module.bias.data.zero_() |
| elif isinstance(module, nn.Embedding): |
| module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) |
| if module.padding_idx is not None: |
| module.weight.data[module.padding_idx].zero_() |
| elif isinstance(module, nn.LayerNorm): |
| module.bias.data.zero_() |
| module.weight.data.fill_(1.0) |
| def get_input_embeddings(self): |
| return self.embedding_layer |
| def set_input_embeddings(self, value): |
| self.embedding_layer = value |
| def get_output_embeddings(self): |
| return self.language_model_head |
| def set_output_embeddings(self, new_embeddings): |
| self.language_model_head = new_embeddings |
| def tie_weights(self): |
| self.language_model_head.weight = self.embedding_layer.weight |
| def get_attention_mask(self, input_ids): |
| batch_size, seq_len = input_ids.shape |
| device = input_ids.device |
| causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool() |
| attention_mask = torch.zeros(batch_size, 1, seq_len, seq_len, device=device) |
| attention_mask[:, :, causal_mask] = float('-inf') |
| padding_mask = (input_ids == self.config.pad_token_id) |
| valid_mask = ~padding_mask |
| if padding_mask.any(): |
| padding_mask_expanded = padding_mask.unsqueeze(1).unsqueeze(2) |
| attention_mask = attention_mask.masked_fill(padding_mask_expanded, float('-inf')) |
| return attention_mask, valid_mask |
| def forward(self, input_ids, labels=None, attention_mask=None, **kwargs): |
| batch_size, seq_len = input_ids.shape |
| hidden_states = self.embedding_layer(input_ids) |
| attn_mask, padding_mask = self.get_attention_mask(input_ids) |
| for layer in self.base_transformer: |
| hidden_states = layer(hidden_states, attn_mask) |
| complexity_logits, complexity_class, recursion_steps = self.router( |
| hidden_states, padding_mask |
| ) |
| if self.training: |
| max_steps = self.config.complex_recursion_steps |
| for step in range(max_steps): |
| hidden_states = self.recursion_layer(hidden_states, attn_mask) |
| else: |
| max_steps_in_batch = int(recursion_steps.max().item()) |
| for step in range(max_steps_in_batch): |
| step_mask = (recursion_steps > step).float().unsqueeze(-1).unsqueeze(-1) |
| new_hidden = self.recursion_layer(hidden_states, attn_mask) |
| hidden_states = step_mask * new_hidden + (1 - step_mask) * hidden_states |
| hidden_states = self.final_layer_norm(hidden_states) |
| logits = self.language_model_head(hidden_states) |
| loss = None |
| if labels is not None: |
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
| loss_fct = nn.CrossEntropyLoss(ignore_index=-100) |
| lm_loss = loss_fct( |
| shift_logits.view(-1, self.config.vocab_size), |
| shift_labels.view(-1) |
| ) |
| complexity_value = min(max(seq_len // 170, 0), 2) |
| pseudo_labels = torch.full( |
| (batch_size,), |
| complexity_value, |
| dtype=torch.long, |
| device=input_ids.device |
| ) |
| router_loss_fct = nn.CrossEntropyLoss() |
| router_loss = router_loss_fct(complexity_logits, pseudo_labels) |
| loss = lm_loss + 0.1 * router_loss |
| return CausalLMOutputWithPast( |
| loss=loss, |
| logits=logits, |
| ) |
| def generate(self, input_ids, max_new_tokens=50, temperature=1.0, |
| top_p=0.9, do_sample=True, **kwargs): |
| self.eval() |
| generated = input_ids |
| for _ in range(max_new_tokens): |
| with torch.no_grad(): |
| outputs = self.forward(generated) |
| logits = outputs.logits |
| next_token_logits = logits[:, -1, :] / temperature |
| if do_sample: |
| sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True) |
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) |
| sorted_indices_to_remove = cumulative_probs > top_p |
| sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() |
| sorted_indices_to_remove[..., 0] = 0 |
| indices_to_remove = sorted_indices_to_remove.scatter( |
| 1, sorted_indices, sorted_indices_to_remove |
| ) |
| next_token_logits[indices_to_remove] = float('-inf') |
| probs = F.softmax(next_token_logits, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
| else: |
| next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) |
| generated = torch.cat([generated, next_token], dim=-1) |
| if next_token.item() == self.config.eos_token_id: |
| break |
| return generated |
