Update app.py

app.py

@@ -6,7 +6,7 @@ from torch.utils.data import Dataset, DataLoader
 import os
 import re
 import time 
-import torch.nn.functional as F # <<<<<<<<<<<< ADDED THIS IMPORT
+import torch.nn.functional as F
 from model import SWCKModel, SeedParser, EntropyEstimator 
 
 # --- Vocabulary and Tokenizer Setup ---
@@ -57,6 +57,17 @@ OVERALL_OUTPUT_ENTROPY_REG_WEIGHT_APP = 0.01
 GATE_SPARSITY_LOSS_WEIGHT_APP = 0.001
 WIRING_PHASE_EPOCHS_APP = 1 
 
+# --- Helper to toggle all debug prints in the model ---
+def set_model_debug_prints(model, seed_parser_debug, block_debug, model_debug):
+    if model:
+        model.debug_prints_enabled = model_debug
+        if hasattr(model, 'seed_parser'):
+            model.seed_parser.debug_prints_enabled = seed_parser_debug
+        if hasattr(model, 'adaptive_blocks'):
+            for block in model.adaptive_blocks:
+                block.debug_prints_enabled = block_debug
+        print(f"App: Model debug prints set - SeedParser: {seed_parser_debug}, Blocks: {block_debug}, SWCKModel: {model_debug}")
+
 
 def build_vocab_from_corpus_text_app(corpus_text):
     global VOCAB_SIZE_APP
@@ -74,7 +85,7 @@ def build_vocab_from_corpus_text_app(corpus_text):
     print(f"App: Built vocab of size {VOCAB_SIZE_APP}")
     return temp_word_to_idx, temp_idx_to_word
 
-def initialize_or_load_model_app():
+def initialize_or_load_model_app(enable_initial_debug=True): # Control initial debug prints
     global swck_model_global, optimizer_global, word_to_idx_global, idx_to_word_global, \
            VOCAB_SIZE_APP, model_load_status_global
 
@@ -92,13 +103,16 @@ def initialize_or_load_model_app():
         'seed_number_str': SEED_NUMBER_STR_APP,
         'num_sub_modules_per_block': NUM_SUB_MODULES_PER_BLOCK_APP
     }
-
+    
+    # Temporarily disable debug during model init to avoid clutter if enable_initial_debug is False
+    # The SeedParser within SWCKModel will print if its own flag is True
+    
     swck_model_global = SWCKModel(**model_args).to(device_global)
-    swck_model_global.debug_prints_enabled = True 
-    if hasattr(swck_model_global, 'seed_parser'): swck_model_global.seed_parser.debug_prints_enabled = True
-    for i,block in enumerate(swck_model_global.adaptive_blocks):
-        block.debug_prints_enabled = True 
-        # print(f"App: Debug prints explicitly enabled for AdaptiveBlock {i}")
+    # Set debug prints AFTER full model initialization
+    set_model_debug_prints(swck_model_global, 
+                           seed_parser_debug=enable_initial_debug, 
+                           block_debug=enable_initial_debug, 
+                           model_debug=enable_initial_debug)
 
 
     if os.path.exists(CHECKPOINT_FILENAME):
@@ -126,8 +140,13 @@ def initialize_or_load_model_app():
             model_load_status_global = f"Model loaded successfully from {CHECKPOINT_FILENAME}."
             print(model_load_status_global)
         except Exception as e:
-            print(f"App: Error loading model from checkpoint: {e}. Initializing new model.")
-            swck_model_global = SWCKModel(**model_args).to(device_global) # Re-init
+            print(f"App: Error loading model from checkpoint: {e}. Re-initializing new model.")
+            # Re-initialize if loading failed, ensuring debug flags are set again
+            swck_model_global = SWCKModel(**model_args).to(device_global)
+            set_model_debug_prints(swck_model_global, 
+                                   seed_parser_debug=enable_initial_debug, 
+                                   block_debug=enable_initial_debug, 
+                                   model_debug=enable_initial_debug)
             optimizer_global = optim.AdamW(swck_model_global.parameters(), lr=0.001)
             model_load_status_global = "Error loading checkpoint. Using new (untrained) model."
     else:
@@ -135,7 +154,13 @@ def initialize_or_load_model_app():
         optimizer_global = optim.AdamW(swck_model_global.parameters(), lr=0.001)
         model_load_status_global = "Initialized a new (untrained) model."
     
-    swck_model_global.eval() 
+    swck_model_global.eval() # Default to eval mode
+    # After loading or initializing, ensure debug prints are set based on desire for startup logs
+    # If enable_initial_debug was False, they are off. If True, they were on during init.
+    # For operations like training/generation, we'll toggle them explicitly.
+    if not enable_initial_debug: # Turn them off if they weren't meant to be on for init
+        set_model_debug_prints(swck_model_global, False, False, False)
+
     return model_load_status_global
 
 
@@ -195,17 +220,16 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
         swck_model_global.set_wiring_phase(epoch < WIRING_PHASE_EPOCHS_APP) 
         epoch_loss = 0.0
         
-        first_batch_debug = (epoch == 0)
+        # Enable full debug for the first batch of the first "wiring" epoch
+        # This will give detailed insight into the "self-wiring roll" on the first piece of data
+        is_first_wiring_batch = (epoch < WIRING_PHASE_EPOCHS_APP and epoch == 0)
 
         for batch_idx, (src_batch, tgt_batch) in enumerate(app_dataloader):
-            if first_batch_debug and batch_idx == 0:
-                swck_model_global.debug_prints_enabled = True
-                if hasattr(swck_model_global, 'seed_parser'): swck_model_global.seed_parser.debug_prints_enabled = True
-                for blk in swck_model_global.adaptive_blocks: blk.debug_prints_enabled = True
-            elif not (first_batch_debug and batch_idx == 0) : 
-                swck_model_global.debug_prints_enabled = False
-                if hasattr(swck_model_global, 'seed_parser'): swck_model_global.seed_parser.debug_prints_enabled = False
-                for blk in swck_model_global.adaptive_blocks: blk.debug_prints_enabled = False
+            if is_first_wiring_batch and batch_idx == 0:
+                print(">>> Enabling FULL DEBUG for first wiring batch <<<")
+                set_model_debug_prints(swck_model_global, True, True, True)
+            else: # Otherwise, keep debug prints minimal or off for speed
+                set_model_debug_prints(swck_model_global, False, False, False)
 
 
             src_batch, tgt_batch = src_batch.to(device_global), tgt_batch.to(device_global)
@@ -231,7 +255,7 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
             if entropy_report["block_output_entropies"]:
                 for i, block_entropy_tensor in enumerate(entropy_report["block_output_entropies"]):
                     target_entropy_val = swck_model_global.seed_parser.get_block_config(i)["target_entropy"]
-                    block_entropy_loss += F.mse_loss(block_entropy_tensor, torch.tensor(target_entropy_val, device=device_global)) # Used F here
+                    block_entropy_loss += F.mse_loss(block_entropy_tensor, torch.tensor(target_entropy_val, device=device_global)) 
                 if entropy_report["block_output_entropies"]: 
                     block_entropy_loss = block_entropy_loss / len(entropy_report["block_output_entropies"])
 
@@ -254,23 +278,19 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
             epoch_loss += combined_loss.item()
 
             log_line = f"  Epoch {epoch+1}, Batch {batch_idx+1}/{len(app_dataloader)}, Loss: {combined_loss.item():.4f}"
-            if batch_idx % max(1, len(app_dataloader)//5) == 0 or batch_idx == len(app_dataloader)-1 : 
+            if batch_idx % max(1, len(app_dataloader)//2) == 0 or batch_idx == len(app_dataloader)-1 : 
                 print(log_line) 
                 training_log_output += log_line + "\n"
         
-        swck_model_global.debug_prints_enabled = False
-        if hasattr(swck_model_global, 'seed_parser'): swck_model_global.seed_parser.debug_prints_enabled = False
-        for blk in swck_model_global.adaptive_blocks: blk.debug_prints_enabled = False
-
+        # Ensure debug is off after the first special batch
+        set_model_debug_prints(swck_model_global, False, False, False)
 
         avg_epoch_loss = epoch_loss / len(app_dataloader) if len(app_dataloader) > 0 else epoch_loss
         epoch_summary = f"Epoch {epoch+1}/{num_epochs_app} - Avg Loss: {avg_epoch_loss:.4f}\n"
         print(epoch_summary)
         training_log_output += epoch_summary
     
-    swck_model_global.debug_prints_enabled = False
-    if hasattr(swck_model_global, 'seed_parser'): swck_model_global.seed_parser.debug_prints_enabled = False
-    for blk in swck_model_global.adaptive_blocks: blk.debug_prints_enabled = False
+    set_model_debug_prints(swck_model_global, False, False, False) # Ensure off after all training
     swck_model_global.eval() 
     
     try:
@@ -296,7 +316,7 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
 
     return training_log_output
 
-def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
+def generate_text_for_app(prompt_str, max_len_gen, temperature_gen, enable_gen_debug: bool): # Add debug toggle
     global model_load_status_global 
     if swck_model_global is None or word_to_idx_global is None or idx_to_word_global is None:
         return "Model not loaded. Please check server logs or try training.", "Model not available."
@@ -304,7 +324,10 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
     swck_model_global.eval() 
     swck_model_global.set_wiring_phase(False) 
     
-    print(f"App: Generating for prompt: '{prompt_str}', max_len: {max_len_gen}, temp: {temperature_gen}")
+    # Set debug prints based on UI toggle for this generation call
+    set_model_debug_prints(swck_model_global, enable_gen_debug, enable_gen_debug, enable_gen_debug)
+    
+    print(f"App: Generating for prompt: '{prompt_str}', max_len: {max_len_gen}, temp: {temperature_gen}, Debug: {enable_gen_debug}")
 
     tokens = [SOS_TOKEN] + [word_to_idx_global.get(w, UNK_TOKEN) for w in prompt_str.lower().split()]
     generated_ids_app = list(tokens)
@@ -324,7 +347,7 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
             if temperature_gen == 0: 
                 next_token_id = torch.argmax(next_token_logits).item()
             else:
-                probs = F.softmax(next_token_logits / temperature_gen, dim=-1) # Used F here
+                probs = F.softmax(next_token_logits / temperature_gen, dim=-1) 
                 if probs.isnan().any() or probs.isinf().any() or torch.sum(probs).item() < 1e-9 : 
                     print(f"Warning: Invalid probabilities at step {i}. Using uniform.")
                     probs = torch.ones_like(next_token_logits) / next_token_logits.size(-1) 
@@ -335,7 +358,7 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
                 break
             generated_ids_app.append(next_token_id)
             
-            if i < 10 : 
+            if i < 10 : # UI debug info is still limited to first 10 new tokens for brevity
                 current_word = idx_to_word_global.get(next_token_id, UNK_TOKEN_STR)
                 overall_ent = entropy_report_infer['overall_output_entropy'].item()
                 if entropy_report_infer['block_output_entropies'] and len(entropy_report_infer['block_output_entropies']) > 0:
@@ -357,9 +380,12 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
 
     debug_output_str = "\n".join(debug_info_lines)
     
+    # Important: Turn off debug prints after generation if they were turned on
+    set_model_debug_prints(swck_model_global, False, False, False)
     return final_text, debug_output_str
 
-initial_load_status = initialize_or_load_model_app()
+# Load model once on app startup. Set enable_initial_debug=False for cleaner startup logs.
+initial_load_status = initialize_or_load_model_app(enable_initial_debug=False) 
 
 with gr.Blocks(title="SWCK Conceptual Demo") as demo:
     model_status_md = gr.Markdown(value=f"**Model Status:** {initial_load_status}", elem_id="model_status_md_123")
@@ -375,30 +401,32 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
         with gr.TabItem("Generate Text"):
             with gr.Row():
                 prompt_input = gr.Textbox(label="Enter your prompt:", placeholder="e.g., the meaning of existence is", scale=3)
+                enable_generation_debug_checkbox = gr.Checkbox(label="Enable Full Kernel Debug (to Console Logs)", value=False)
+            with gr.Row():
                 generate_button = gr.Button("Generate", scale=1)
             with gr.Row():
                 max_len_slider = gr.Slider(minimum=10, maximum=150, value=50, step=1, label="Max Generation Length")
                 temp_slider = gr.Slider(minimum=0.0, maximum=2.0, value=0.8, step=0.1, label="Temperature (0 for greedy)")
             
             output_text = gr.Textbox(label="Generated Text:", lines=6, interactive=False)
-            debug_text_area = gr.Textbox(label="Generation Debug Info (first few steps):", lines=8, interactive=False)
+            debug_text_area = gr.Textbox(label="Generation Debug Info (first few steps to UI):", lines=8, interactive=False)
 
         with gr.TabItem("In-App Training (Conceptual Test)"):
-            gr.Markdown("WARNING: In-app training is EXTREMELY slow and only for basic conceptual testing on Spaces free tier. Uses a small internal corpus. Model state persists only for this session unless saved manually via code modification.")
+            gr.Markdown("WARNING: In-app training is EXTREMELY slow and only for basic conceptual testing on Spaces free tier. Uses a small internal corpus. Model state persists only for this session unless saved manually via code modification. Full Kernel Debug will be printed to console for the FIRST BATCH of the FIRST WIRING EPOCH ONLY.")
             with gr.Row():
-                train_epochs_slider = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of Training Epochs")
-                train_batch_size_slider = gr.Slider(minimum=1, maximum=8, value=2, step=1, label="Training Batch Size")
+                train_epochs_slider = gr.Slider(minimum=1, maximum=3, value=1, step=1, label="Number of Training Epochs (1-3 for demo)") # Reduced max
+                train_batch_size_slider = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Training Batch Size (1-4 for demo)") # Reduced max
                 train_lr_slider = gr.Slider(minimum=1e-5, maximum=1e-3, value=5e-4, step=1e-5, label="Learning Rate") 
             
             start_training_button = gr.Button("Start Short Training Session")
-            training_status_output = gr.Textbox(label="Training Log / Status:", lines=10, interactive=False,show_label=True )
+            training_status_output = gr.Textbox(label="Training Log / Status (summary):", lines=10, interactive=False,show_label=True )
 
     def update_status_text_for_ui(): 
         return f"**Model Status:** {model_load_status_global}"
 
     generate_button.click(
         fn=generate_text_for_app,
-        inputs=[prompt_input, max_len_slider, temp_slider],
+        inputs=[prompt_input, max_len_slider, temp_slider, enable_generation_debug_checkbox], # Added checkbox
         outputs=[output_text, debug_text_area]
     )
     
@@ -410,4 +438,6 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
     
 
 if __name__ == "__main__":
+    # For local testing, you can launch with debug=True for Gradio's server debug.
+    # The model's internal debug prints are controlled by set_model_debug_prints().
     demo.launch(debug=True)