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=""): super().__init__() self.fc1 = nn.Linear(d_model, hidden_dim) self.fc2 = nn.Linear(hidden_dim, 1) self.name = name self.debug_prints_enabled = True # Default to True for this module if needed def forward(self, x, active_mask=None): # x: (batch, seq_len, d_model) # Simplified masking logic for robustness if x.numel() == 0: return torch.tensor(0.0, device=x.device) if active_mask is not None: # Ensure active_mask is boolean and compatible shape for broadcasting/indexing if active_mask.dtype != torch.bool: active_mask = active_mask.bool() if x.dim() == 3 and active_mask.dim() == 2 and x.shape[:2] == active_mask.shape: # typical case: x is (B,S,D), active_mask is (B,S) x_masked = x[active_mask] # This flattens to (N_active, D) elif x.dim() == 2 and active_mask.dim() == 1 and x.shape[0] == active_mask.shape[0]: # x is (S,D) or (B,D) - less common here, but handle x_masked = x[active_mask] else: # Fallback if mask shapes are unexpected, process all elements # if self.debug_prints_enabled: # print(f"Warning [{self.name}]: Mask shape mismatch (x: {x.shape}, mask: {active_mask.shape}). Processing all elements.") x_masked = x.reshape(-1, x.size(-1)) else: x_masked = x.reshape(-1, x.size(-1)) if x_masked.numel() == 0: return torch.tensor(0.0, device=x.device) h = F.relu(self.fc1(x_masked)) # Sigmoid output, then mean. Represents average "activity" or "confidence" as a proxy for entropy. estimated_entropy = torch.sigmoid(self.fc2(h)).mean() return estimated_entropy # --- 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 if self.debug_prints_enabled: print(f"--- SeedParser Initialization ---") print(f" Seed Phrase (start): '{self.seed_phrase[:50]}...'") print(f" Seed Number: {self.seed_number_str}") phrase_hash = hashlib.sha256(seed_phrase.encode()).hexdigest() self.phrase_base_val = int(phrase_hash[:16], 16) if self.debug_prints_enabled: print(f" Phrase Base Value (from hash): {self.phrase_base_val}") self.num_sequence = [int(d) for d in seed_number_str if d.isdigit()] if not self.num_sequence: self.num_sequence = [sum(bytearray(seed_number_str.encode())) % 10] 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" SeedParser: Generated InitMap:") for i, block_config in enumerate(self.init_map["block_configs"]): gate_inits_str = [f'{g:.3f}' for g in block_config['initial_gate_proportions']] print(f" Block {i}: Target Entropy: {block_config['target_entropy']:.4f}, Initial Gate Proportions: {gate_inits_str}") if self.debug_prints_enabled: print(f"--- SeedParser Initialized ---") def _get_deterministic_value(self, key_name, min_val, max_val, sequence_idx_offset=0): key_specific_hash = int(hashlib.sha256(key_name.encode() + self.seed_phrase.encode()).hexdigest()[:8], 16) num_seq_val = 0 if self.num_sequence: for i, digit in enumerate(self.num_sequence): num_seq_val = (num_seq_val * 10 + digit) % 1000003 combined_seed_val = self.phrase_base_val + key_specific_hash + num_seq_val + sequence_idx_offset if max_val == min_val: return min_val val_range = max_val - min_val + 1 return min_val + int(abs(math.sin(float(combined_seed_val)) * 1e5)) % val_range def _get_deterministic_float(self, key_name, min_val=0.0, max_val=1.0, sequence_idx_offset=0): key_specific_hash = int(hashlib.sha256(key_name.encode() + self.seed_phrase.encode()).hexdigest()[:8], 16) num_seq_val = 0 if self.num_sequence: for i, digit in enumerate(self.num_sequence): num_seq_val = (num_seq_val * 10 + digit) % 1000003 combined_seed_val = self.phrase_base_val + key_specific_hash + num_seq_val + sequence_idx_offset norm_float = (math.sin(float(combined_seed_val) * 0.1) + 1.0) / 2.0 scaled_val = min_val + norm_float * (max_val - min_val) return scaled_val def _generate_init_map(self): init_map = {"block_configs": []} for i in range(self.num_adaptive_blocks): gate_raw_scores = [ self._get_deterministic_float(f"block_{i}_gate_{j}_raw_score", -1.0, 1.0, sequence_idx_offset=i*10 + j) for j in range(self.num_sub_modules_per_block) ] if self.num_sub_modules_per_block > 0: gate_initial_proportions = F.softmax(torch.tensor(gate_raw_scores), dim=0).tolist() else: gate_initial_proportions = [] target_entropy = self._get_deterministic_float( f"block_{i}_target_entropy", 0.05, 0.35, sequence_idx_offset=i ) init_map["block_configs"].append({ "initial_gate_proportions": gate_initial_proportions, "raw_gate_scores_for_param_init": gate_raw_scores, "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_for_block, 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_for_block self.debug_prints_enabled = True if self.debug_prints_enabled: print(f" Initializing AdaptiveBlock {self.block_idx} with seed config: TargetEntropy={self.config_from_seed['target_entropy']:.3f}, InitialGateProportions={[f'{g:.3f}' for g in self.config_from_seed['initial_gate_proportions']]}") 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)) 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)) 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)} defined. Using defined count.") self.num_sub_modules = len(self.sub_modules) raw_gate_param_inits = self.config_from_seed.get("raw_gate_scores_for_param_init", [0.0] * self.num_sub_modules if self.num_sub_modules > 0 else []) if len(raw_gate_param_inits) != self.num_sub_modules: print(f"Warning: Block {self.block_idx} raw_gate_scores length mismatch. Re-initializing to zeros.") raw_gate_param_inits = [0.0] * self.num_sub_modules if self.num_sub_modules > 0 else [] self.gates_params = nn.Parameter(torch.tensor(raw_gate_param_inits, dtype=torch.float32)) self.initial_gate_proportions_tensor = torch.tensor(self.config_from_seed['initial_gate_proportions'], dtype=torch.float32) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) self.output_entropy_estimator = EntropyEstimator(d_model, name=f"Block{block_idx}_OutEntropy") self.wiring_phase_active = False def set_wiring_phase(self, active): self.wiring_phase_active = active # if self.debug_prints_enabled: # phase_status = "ACTIVATED" if active else "DEACTIVATED" # print(f" AdaptiveBlock {self.block_idx}: WIRING PHASE {phase_status}") # Made less verbose def forward(self, x, key_padding_mask=None, attn_mask=None): current_gates_softmax = F.softmax(self.gates_params, dim=0) # if self.debug_prints_enabled: # Made less verbose # print(f" AdaptiveBlock {self.block_idx} Input x: {x.shape}, Current Gates (softmax): {[f'{g.item():.3f}' for g in current_gates_softmax]}") x_norm = self.norm1(x) outputs = [] for i, module in enumerate(self.sub_modules): if i >= self.num_sub_modules: break if i == 0: module_out, _ = module(x_norm, x_norm, x_norm, key_padding_mask=key_padding_mask, attn_mask=attn_mask, need_weights=False) else: module_out = module(x_norm) outputs.append(module_out) if not outputs: if self.debug_prints_enabled: print(f" AdaptiveBlock {self.block_idx}: No sub_modules processed. Passing input through.") final_out_unnorm = x else: stacked_outputs = torch.stack(outputs, dim=0) weighted_sum = torch.sum(stacked_outputs * current_gates_softmax.view(-1, 1, 1, 1), dim=0) final_out_unnorm = x + self.dropout(weighted_sum) final_out_norm = self.norm2(final_out_unnorm) 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) if self.wiring_phase_active and self.training: with torch.no_grad(): entropy_diff = current_output_entropy - target_entropy_for_block adjustment_strength = 0.01 if entropy_diff > 0.05: self.gates_params.data[1] += adjustment_strength if self.num_sub_modules > 2: self.gates_params.data[2] += adjustment_strength self.gates_params.data[0] -= adjustment_strength * 0.5 elif entropy_diff < -0.05: self.gates_params.data[0] += adjustment_strength self.gates_params.data[1] -= adjustment_strength * 0.5 if self.num_sub_modules > 2: self.gates_params.data[2] -= adjustment_strength * 0.5 self.gates_params.data.clamp_(-2.5, 2.5) 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 Gate Params (raw): {[f'{g.item():.3f}' for g in self.gates_params.data]}") initial_gate_targets_on_device = self.initial_gate_proportions_tensor.to(self.gates_params.device) return final_out_norm, current_output_entropy, current_gates_softmax, self.gates_params, initial_gate_targets_on_device # --- Positional Encoding --- class PositionalEncoding(nn.Module): def __init__(self,d_model,dropout=0.1,max_len=512): # Default max_len is good 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)) def forward(self,x): # x: (batch, seq_len, d_model) # self.pe: (1, max_len, d_model) # We need to select the part of pe corresponding to x's seq_len 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 if self.debug_prints_enabled: print(f"--- Initializing SWCKModel ---") self.seed_parser = SeedParser(seed_phrase, seed_number_str, d_model, num_adaptive_blocks, num_sub_modules_per_block) self.seed_parser.debug_prints_enabled = self.debug_prints_enabled self.embedding = nn.Embedding(vocab_size, d_model) # Corrected: PositionalEncoding uses its own default max_len or a hardcoded one. # It does not depend on SEQ_LEN_APP from app.py. 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}") new_block = AdaptiveBlock(d_model, n_heads, d_ff, dropout, block_config, block_idx=i, num_sub_modules=num_sub_modules_per_block) new_block.debug_prints_enabled = self.debug_prints_enabled self.adaptive_blocks.append(new_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.overall_output_entropy_estimator.debug_prints_enabled = self.debug_prints_enabled self._init_weights() if self.debug_prints_enabled: print(f"--- SWCKModel Initialized (Vocab: {vocab_size}, d_model: {d_model}) ---") 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.") # Made less verbose pass for block in self.adaptive_blocks: block.set_wiring_phase(active) def forward(self, src_tokens, src_key_padding_mask=None): # if self.debug_prints_enabled: # Made less verbose # 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} (True means pad)") 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}") # Made less verbose block_output_entropies = [] current_block_gate_softmaxes = [] current_block_gate_params = [] initial_block_gate_targets = [] for i, block in enumerate(self.adaptive_blocks): # if self.debug_prints_enabled: print(f" Processing AdaptiveBlock {i}...") # Made less verbose x, block_entropy, current_gate_softmax, current_gate_param, initial_gate_target = block(x, key_padding_mask=src_key_padding_mask, attn_mask=None) block_output_entropies.append(block_entropy) current_block_gate_softmaxes.append(current_gate_softmax) current_block_gate_params.append(current_gate_param) initial_block_gate_targets.append(initial_gate_target) # if self.debug_prints_enabled: print(f" Output x from AdaptiveBlock {i}: {x.shape}, Entropy: {block_entropy.item():.4f}") # Made less verbose logits = self.fc_out(x) # if self.debug_prints_enabled: print(f" Output logits: {logits.shape}") # Made less verbose 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}") # Made less verbose entropy_report = { "block_output_entropies": block_output_entropies, "overall_output_entropy": overall_entropy, "current_block_gate_softmaxes": current_block_gate_softmaxes, "current_block_gate_params": current_block_gate_params, "initial_block_gate_targets": initial_block_gate_targets } return logits, entropy_report