import torch import torch.nn as nn import torch.nn.functional as F import math import hashlib # For generating deterministic values from seed # --- Helper: Entropy Estimator --- class EntropyEstimator(nn.Module): def __init__(self, d_model, hidden_dim=32, name=""): # Smaller hidden_dim for simplicity super().__init__() self.fc1 = nn.Linear(d_model, hidden_dim) self.fc2 = nn.Linear(hidden_dim, 1) self.name = name def forward(self, x, active_mask=None): # x: (batch, seq_len, d_model) if active_mask is not None and x.shape[:-1] != active_mask.shape: print(f"Warning [{self.name}]: x shape {x.shape[:-1]} and active_mask shape {active_mask.shape} mismatch. Entropy might be inaccurate.") # Fallback if mask is problematic, or process only unmasked if shapes allow if x.numel() == 0: return torch.tensor(0.0, device=x.device) # Handle empty tensor case if active_mask.sum() == 0: return torch.tensor(0.0, device=x.device) # Handle all masked case # Try to apply mask if possible, otherwise average all. This part can be tricky. # For now, if shapes mismatch significantly, we might average all as a robust fallback. # A more robust solution would ensure masks are always correct upstream. if x.dim() == active_mask.dim() + 1 and x.shape[:-1] == active_mask.shape : # (B,S,D) and (B,S) x_masked = x[active_mask] if x_masked.numel() == 0: return torch.tensor(0.0, device=x.device) h = F.relu(self.fc1(x_masked)) return torch.sigmoid(self.fc2(h)).mean() # Mean entropy over active elements else: # Fallback if mask application is uncertain h = F.relu(self.fc1(x.reshape(-1, x.size(-1)))) return torch.sigmoid(self.fc2(h)).mean() elif active_mask is None and x.numel() > 0: h = F.relu(self.fc1(x.reshape(-1, x.size(-1)))) return torch.sigmoid(self.fc2(h)).mean() elif x.numel() == 0: return torch.tensor(0.0, device=x.device) # Handle empty tensor # Default if active_mask is present and correct x_masked = x[active_mask] if x_masked.numel() == 0: return torch.tensor(0.0, device=x.device) h = F.relu(self.fc1(x_masked)) return torch.sigmoid(self.fc2(h)).mean() # Mean entropy over active elements # --- Helper: Seed Parser --- class SeedParser: def __init__(self, seed_phrase, seed_number_str, d_model, num_adaptive_blocks, num_sub_modules_per_block): self.seed_phrase = seed_phrase self.seed_number_str = seed_number_str self.d_model = d_model self.num_adaptive_blocks = num_adaptive_blocks self.num_sub_modules_per_block = num_sub_modules_per_block self.debug_prints_enabled = True print(f"--- SeedParser Initialization ---") print(f" Seed Phrase: '{self.seed_phrase}'") print(f" Seed Number: {self.seed_number_str}") # 1. Process Seed Phrase (e.g., to get a base vector) # For simplicity, hash it to get a deterministic starting point for numerical derivation phrase_hash = hashlib.sha256(seed_phrase.encode()).hexdigest() self.phrase_base_val = int(phrase_hash[:8], 16) # Use first 8 hex chars if self.debug_prints_enabled: print(f" Phrase Base Value (from hash): {self.phrase_base_val}") # 2. Process Seed Number (more direct influence on structure) self.num_sequence = [int(d) for d in seed_number_str if d.isdigit()] if not self.num_sequence: self.num_sequence = [0] # Fallback if self.debug_prints_enabled: print(f" Numerical Sequence (from seed number): {self.num_sequence}") self.init_map = self._generate_init_map() if self.debug_prints_enabled: print(f" Generated InitMap:") for i, block_config in enumerate(self.init_map["block_configs"]): print(f" Block {i}: Active Module Index: {block_config['active_module_idx']}, Target Entropy: {block_config['target_entropy']:.4f}, Gate Inits: {[f'{g:.2f}' for g in block_config['gate_inits']]}") print(f"--- SeedParser Initialized ---") def _get_deterministic_value(self, key_name, min_val, max_val, sequence_idx_offset=0): # Combine phrase base and numerical sequence for more variation combined_seed_val = self.phrase_base_val for i, num in enumerate(self.num_sequence): combined_seed_val += num * (10**(i + sequence_idx_offset)) # Hash the key_name to make it specific to the parameter key_hash = int(hashlib.sha256(key_name.encode()).hexdigest()[:8], 16) final_seed = combined_seed_val + key_hash # Simple mapping to range (not cryptographically strong, but deterministic) if max_val == min_val: return min_val # Avoid division by zero if range is 1 val = min_val + (final_seed % (max_val - min_val + 1)) return val def _get_deterministic_float(self, key_name, min_val=0.0, max_val=1.0, sequence_idx_offset=0): combined_seed_val = self.phrase_base_val for i, num in enumerate(self.num_sequence): combined_seed_val += num * (10**(i + sequence_idx_offset)) key_hash = int(hashlib.sha256(key_name.encode()).hexdigest()[:8], 16) final_seed = combined_seed_val + key_hash # Map to [0,1] float then scale float_val = (final_seed % 1000001) / 1000000.0 # Ensure it's never exactly 0 for some ops scaled_val = min_val + float_val * (max_val - min_val) return scaled_val def _generate_init_map(self): init_map = {"block_configs": []} for i in range(self.num_adaptive_blocks): # Determine which sub-module is initially "more" active active_module_idx = self._get_deterministic_value( f"block_{i}_active_module", 0, self.num_sub_modules_per_block - 1, sequence_idx_offset=i ) # Determine initial gating values (summing to 1 for softmax-like behavior later) gate_inits_raw = [ self._get_deterministic_float(f"block_{i}_gate_{j}_init_raw", 0.1, 1.0, sequence_idx_offset=i*10 + j) for j in range(self.num_sub_modules_per_block) ] # Make one gate stronger based on active_module_idx, then normalize slightly if self.num_sub_modules_per_block > 0 : gate_inits_raw[active_module_idx] *= 2.0 # Boost the 'active' one sum_raw = sum(gate_inits_raw) gate_inits_normalized = [g / sum_raw for g in gate_inits_raw] if sum_raw > 0 else [1.0/self.num_sub_modules_per_block]*self.num_sub_modules_per_block else: gate_inits_normalized = [] # Determine a target entropy for this block's output target_entropy = self._get_deterministic_float( f"block_{i}_target_entropy", 0.05, 0.3, sequence_idx_offset=i # Target a moderate, non-zero entropy ) init_map["block_configs"].append({ "active_module_idx": active_module_idx, # For initial bias "gate_inits": gate_inits_normalized, # Initial values for learnable gates "target_entropy": target_entropy }) return init_map def get_block_config(self, block_idx): if 0 <= block_idx < len(self.init_map["block_configs"]): return self.init_map["block_configs"][block_idx] return None # --- Adaptive Block --- class AdaptiveBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout, seed_parser_config, block_idx, num_sub_modules=3): super().__init__() self.d_model = d_model self.block_idx = block_idx self.num_sub_modules = num_sub_modules self.config_from_seed = seed_parser_config # dict for this block self.debug_prints_enabled = True if self.debug_prints_enabled: print(f" Initializing AdaptiveBlock {self.block_idx} with seed config: {self.config_from_seed}") # Define potential sub-modules self.sub_module_0 = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) self.sub_module_1 = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) # Sub-module 2: A simpler FFN or even a near identity (residual + small transform) self.sub_module_2 = nn.Sequential( nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model // 2, d_model) ) # Add more diverse sub-modules if needed for `num_sub_modules_per_block` self.sub_modules = nn.ModuleList([self.sub_module_0, self.sub_module_1, self.sub_module_2]) if self.num_sub_modules > len(self.sub_modules): print(f"Warning: block {self.block_idx} requested {self.num_sub_modules} sub_modules, but only {len(self.sub_modules)} are defined. Using defined ones.") self.num_sub_modules = len(self.sub_modules) # Learnable gates for combining/selecting sub-modules # Initialize gates based on seed_parser_config gate_initial_values = self.config_from_seed.get("gate_inits", [1.0/self.num_sub_modules]*self.num_sub_modules if self.num_sub_modules > 0 else []) if len(gate_initial_values) != self.num_sub_modules: # Fallback if seed parser gave wrong number print(f"Warning: Block {self.block_idx} gate_inits length mismatch. Re-initializing uniformly.") gate_initial_values = [1.0/self.num_sub_modules]*self.num_sub_modules if self.num_sub_modules > 0 else [] self.gates = nn.Parameter(torch.tensor(gate_initial_values, dtype=torch.float32)) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) # For output of block self.dropout = nn.Dropout(dropout) self.output_entropy_estimator = EntropyEstimator(d_model, name=f"Block{block_idx}_OutEntropy") self.wiring_phase_active = False # To be set by the main model def set_wiring_phase(self, active): self.wiring_phase_active = active if self.debug_prints_enabled and active: print(f" AdaptiveBlock {self.block_idx}: WIRING PHASE ACTIVATED") elif self.debug_prints_enabled and not active: print(f" AdaptiveBlock {self.block_idx}: WIRING PHASE DEACTIVATED") def forward(self, x, key_padding_mask=None, attn_mask=None): # attn_mask is for MHA, key_padding_mask for MHA keys if self.debug_prints_enabled: current_gates_softmax = F.softmax(self.gates, dim=0) print(f" AdaptiveBlock {self.block_idx} Input x: {x.shape}, Gates (softmax): {[f'{g.item():.3f}' for g in current_gates_softmax]}") x_norm = self.norm1(x) outputs = [] active_module_found = False for i, module in enumerate(self.sub_modules): if i >= self.num_sub_modules: break # Only use configured number if i == 0: # MHA # MHA expects key_padding_mask (N, S) bool: True if padded. # attn_mask (L,S) or (N*H,L,S) float/bool: True if masked / -inf. # For self-attention, L=S. If attn_mask is causal (L,L), it's fine. # If key_padding_mask is (N,S), it's fine. module_out, _ = module(x_norm, x_norm, x_norm, key_padding_mask=key_padding_mask, attn_mask=attn_mask, need_weights=False) # Don't need weights for this sim active_module_found = True elif hasattr(module, 'fc1') or isinstance(module, nn.Sequential): # FFN-like module_out = module(x_norm) active_module_found = True else: # Fallback for undefined module types in this simple sketch module_out = x_norm # Pass through outputs.append(module_out) if not active_module_found or not outputs: # Should not happen if num_sub_modules > 0 print(f" AdaptiveBlock {self.block_idx}: No active sub_modules processed. Passing input through.") final_out_unnorm = x # pass through else: # Gated combination gate_weights = F.softmax(self.gates, dim=0) # Ensure they sum to 1 # Weighted sum of module outputs # Ensure outputs are stackable (they should be if all modules output (B,S,D)) if outputs: stacked_outputs = torch.stack(outputs, dim=0) # (num_sub_modules, B, S, D) # gate_weights (num_sub_modules) -> (num_sub_modules, 1, 1, 1) for broadcasting weighted_sum = torch.sum(stacked_outputs * gate_weights.view(-1, 1, 1, 1), dim=0) final_out_unnorm = x + self.dropout(weighted_sum) # Residual connection else: # Fallback if somehow no outputs final_out_unnorm = x final_out_norm = self.norm2(final_out_unnorm) # During wiring phase, we might adjust gates based on local entropy vs target # This is a very simplified "self-wiring" heuristic current_output_entropy = self.output_entropy_estimator(final_out_norm, active_mask=~key_padding_mask if key_padding_mask is not None else None) target_entropy_for_block = self.config_from_seed.get("target_entropy", 0.1) # Default target if self.wiring_phase_active and self.training : # Only adjust gates during wiring AND training with torch.no_grad(): # Don't track gradients for this heuristic adjustment entropy_diff = current_output_entropy - target_entropy_for_block # If current entropy is too high, slightly boost gates of modules that might reduce it (heuristic) # If too low, slightly boost gates of modules that might increase it (heuristic) # This is extremely heuristic. A true self-wiring mechanism would be more complex. # For this sketch, let's say MHA (module 0) might increase complexity/entropy if it was low, # and FFNs (module 1, 2) might refine/stabilize if entropy was high. adjustment_strength = 0.01 # Small adjustment if entropy_diff > 0.05: # Current entropy significantly higher than target self.gates.data[1] += adjustment_strength self.gates.data[2] += adjustment_strength self.gates.data[0] -= adjustment_strength * 0.5 # Slightly decrease MHA elif entropy_diff < -0.05: # Current entropy significantly lower self.gates.data[0] += adjustment_strength self.gates.data[1] -= adjustment_strength * 0.5 self.gates.data[2] -= adjustment_strength * 0.5 # Clamp gates to avoid extreme values before softmax (optional) self.gates.data.clamp_(-2.0, 2.0) if self.debug_prints_enabled: print(f" AdaptiveBlock {self.block_idx} WIRING: OutEnt={current_output_entropy.item():.4f}, TgtEnt={target_entropy_for_block:.4f}, Δ={entropy_diff.item():.4f} -> New Gates (raw): {[f'{g.item():.3f}' for g in self.gates.data]}") elif self.debug_prints_enabled: print(f" AdaptiveBlock {self.block_idx} EXEC: OutEnt={current_output_entropy.item():.4f}, TgtEnt={target_entropy_for_block:.4f}") # Return the block's output and its current estimated output entropy return final_out_norm, current_output_entropy, gate_weights # --- Positional Encoding --- class PositionalEncoding(nn.Module): def __init__(self,d_model,dropout=0.1,max_len=512): # Reduced max_len for this sketch super().__init__() self.dropout=nn.Dropout(p=dropout) pe=torch.zeros(max_len,d_model) pos=torch.arange(0,max_len,dtype=torch.float).unsqueeze(1) div=torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model)) pe[:,0::2]=torch.sin(pos*div) pe[:,1::2]=torch.cos(pos*div) self.register_buffer('pe',pe.unsqueeze(0)) # (1, max_len, d_model) def forward(self,x): # x: (batch, seq_len, d_model) x=x+self.pe[:,:x.size(1),:] return self.dropout(x) # --- Main SWCK Model --- class SWCKModel(nn.Module): def __init__(self, vocab_size, d_model, n_heads, d_ff, num_adaptive_blocks, dropout, seed_phrase, seed_number_str, num_sub_modules_per_block=3): super().__init__() self.d_model = d_model self.seed_phrase = seed_phrase self.seed_number_str = seed_number_str self.debug_prints_enabled = True print(f"--- Initializing SWCKModel ---") self.seed_parser = SeedParser(seed_phrase, seed_number_str, d_model, num_adaptive_blocks, num_sub_modules_per_block) self.embedding = nn.Embedding(vocab_size, d_model) self.pos_encoder = PositionalEncoding(d_model, dropout) self.adaptive_blocks = nn.ModuleList() for i in range(num_adaptive_blocks): block_config = self.seed_parser.get_block_config(i) if block_config is None: raise ValueError(f"Could not get seed config for block {i}") self.adaptive_blocks.append( AdaptiveBlock(d_model, n_heads, d_ff, dropout, block_config, block_idx=i, num_sub_modules=num_sub_modules_per_block) ) if self.debug_prints_enabled: print(f" SWCKModel: Added AdaptiveBlock {i}") self.fc_out = nn.Linear(d_model, vocab_size) self.overall_output_entropy_estimator = EntropyEstimator(d_model, name="OverallOutEntropy") self._init_weights() print(f"--- SWCKModel Initialized ---") def _init_weights(self): initrange = 0.1 self.embedding.weight.data.uniform_(-initrange, initrange) self.fc_out.bias.data.zero_() self.fc_out.weight.data.uniform_(-initrange, initrange) def set_wiring_phase(self, active): if self.debug_prints_enabled: print(f"SWCKModel: Setting wiring phase to {active} for all blocks.") for block in self.adaptive_blocks: block.set_wiring_phase(active) def forward(self, src_tokens, src_key_padding_mask=None): # src_tokens: (batch, seq_len) # src_key_padding_mask: (batch, seq_len), True for padded positions if self.debug_prints_enabled: print(f"\n--- SWCKModel Forward Pass ---") print(f" Input src_tokens: {src_tokens.shape}") if src_key_padding_mask is not None: print(f" Input src_key_padding_mask: {src_key_padding_mask.shape}") x = self.embedding(src_tokens) * math.sqrt(self.d_model) x = self.pos_encoder(x) if self.debug_prints_enabled: print(f" After Embedding & PosEnc, x: {x.shape}") block_output_entropies = [] block_gate_weights = [] # For self-attention within blocks, a causal mask might be needed if it's a decoder-style model # For this general "processing core" sketch, let's assume full self-attention unless specified. # If this were a decoder, a causal mask would be passed or generated here. # For now, no explicit top-level causal mask is made, relying on block's internal MHA params. # A more standard transformer would create a causal mask for decoder self-attention. # We'll pass src_key_padding_mask to MHA if it's self-attention on source. for i, block in enumerate(self.adaptive_blocks): if self.debug_prints_enabled: print(f" Processing AdaptiveBlock {i}...") # For self-attention in blocks, key_padding_mask applies to keys/values. # No separate attention mask for now unless it's a decoder block. x, block_entropy, gates = block(x, key_padding_mask=src_key_padding_mask, attn_mask=None) block_output_entropies.append(block_entropy) block_gate_weights.append(gates) if self.debug_prints_enabled: print(f" Output x from AdaptiveBlock {i}: {x.shape}, Entropy: {block_entropy.item():.4f}") logits = self.fc_out(x) if self.debug_prints_enabled: print(f" Output logits: {logits.shape}") # Overall output entropy (of the final representation before fc_out) # Masking for entropy calculation final_active_mask = ~src_key_padding_mask if src_key_padding_mask is not None else None overall_entropy = self.overall_output_entropy_estimator(x, active_mask=final_active_mask) if self.debug_prints_enabled: print(f" Overall Final Representation Entropy: {overall_entropy.item():.4f}") # Entropies from each block, overall output entropy, and gate weights for regularization/logging entropy_report = { "block_output_entropies": block_output_entropies, # List of tensors "overall_output_entropy": overall_entropy, # Tensor "block_gate_weights": block_gate_weights # List of tensors } return logits, entropy_report