Create model.py

model.py

@@ -0,0 +1,390 @@
+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