Update app.py

app.py

@@ -49,7 +49,7 @@ idx_to_word_global = None
 device_global = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_load_status_global = "Model not loaded."
 
-CHECKPOINT_FILENAME = "swck_model_conceptual_app.pth.tar" 
+CHECKPOINT_FILENAME = "swck_model_conceptual_app_fulldebug.pth.tar" # New checkpoint name
 
 MAIN_LOSS_WEIGHT_APP = 1.0
 BLOCK_TARGET_ENTROPY_LOSS_WEIGHT_APP = 0.02
@@ -57,7 +57,6 @@ 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
@@ -85,7 +84,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(enable_initial_debug=True): # Control initial debug prints
+def initialize_or_load_model_app(): 
     global swck_model_global, optimizer_global, word_to_idx_global, idx_to_word_global, \
            VOCAB_SIZE_APP, model_load_status_global
 
@@ -104,15 +103,17 @@ def initialize_or_load_model_app(enable_initial_debug=True): # Control initial d
         '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
-    
+    print("App: Initializing SWCKModel with FULL DEBUG ON by default for init...")
     swck_model_global = SWCKModel(**model_args).to(device_global)
-    # 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)
+    # Debug is on by default in SWCKModel and sub-components as per their class __init__
+    # We can use set_model_debug_prints to confirm or change it if needed later.
+    # For now, rely on their internal defaults being True.
+    # If SeedParser or AdaptiveBlock have their debug_prints_enabled=False by default in model.py,
+    # you would explicitly set them here:
+    if hasattr(swck_model_global, 'seed_parser'): swck_model_global.seed_parser.debug_prints_enabled = True
+    for block in swck_model_global.adaptive_blocks: block.debug_prints_enabled = True
+    swck_model_global.debug_prints_enabled = True
+    print("App: All model component debugs are intended to be ON by default from their init.")
 
 
     if os.path.exists(CHECKPOINT_FILENAME):
@@ -140,27 +141,19 @@ def initialize_or_load_model_app(enable_initial_debug=True): # Control initial d
             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}. Re-initializing new model.")
-            # Re-initialize if loading failed, ensuring debug flags are set again
+            print(f"App: Error loading model from checkpoint: {e}. Re-initializing new model with debug ON.")
             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)
+            if hasattr(swck_model_global, 'seed_parser'): swck_model_global.seed_parser.debug_prints_enabled = True
+            for block in swck_model_global.adaptive_blocks: block.debug_prints_enabled = True
+            swck_model_global.debug_prints_enabled = True
             optimizer_global = optim.AdamW(swck_model_global.parameters(), lr=0.001)
-            model_load_status_global = "Error loading checkpoint. Using new (untrained) model."
+            model_load_status_global = "Error loading checkpoint. Using new (untrained) model with debug ON."
     else:
-        print(f"App: Checkpoint {CHECKPOINT_FILENAME} not found. Initializing new model.")
+        print(f"App: Checkpoint {CHECKPOINT_FILENAME} not found. Initializing new model with debug ON.")
         optimizer_global = optim.AdamW(swck_model_global.parameters(), lr=0.001)
-        model_load_status_global = "Initialized a new (untrained) model."
+        model_load_status_global = "Initialized a new (untrained) model with debug ON."
     
-    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)
-
+    swck_model_global.eval() 
     return model_load_status_global
 
 
@@ -195,12 +188,16 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
     if swck_model_global is None or word_to_idx_global is None:
         return "Model not initialized. Cannot train."
 
-    print("\n--- App: Starting Short Training Session ---")
+    print("\n--- App: Starting Short Training Session (Full Debug ON for ALL batches/epochs) ---")
     progress(0, desc="Preparing training data...")
     
+    # Ensure debug prints are ON for the entire training session
+    set_model_debug_prints(swck_model_global, True, True, True)
+
     training_corpus = SEED_PHRASE_APP + " " + EXTENDED_TEXT_FOR_TRAINING_APP
     app_dataset = AppSWCKDataset(training_corpus, word_to_idx_global, SEQ_LEN_APP, SOS_TOKEN, EOS_TOKEN, PAD_TOKEN)
     if not app_dataset.samples:
+        set_model_debug_prints(swck_model_global, False, False, False) # Turn off if error
         return "App Training Error: No samples created from the corpus."
         
     app_dataloader = DataLoader(app_dataset, batch_size=int(batch_size_app), shuffle=True, collate_fn=app_swck_collate_fn)
@@ -213,24 +210,17 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
 
     criterion_main_app = nn.CrossEntropyLoss(ignore_index=PAD_TOKEN)
     
-    training_log_output = f"Starting training for {num_epochs_app} epochs...\n"
+    training_log_output = f"Starting training for {num_epochs_app} epochs (Full Debug ON)...\n"
     swck_model_global.train() 
 
     for epoch in progress.tqdm(range(int(num_epochs_app)), desc="Training Epochs"):
         swck_model_global.set_wiring_phase(epoch < WIRING_PHASE_EPOCHS_APP) 
         epoch_loss = 0.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)
+        print(f"\n>>> EPOCH {epoch+1} - Starting with Full Debug for all batches <<<")
 
         for batch_idx, (src_batch, tgt_batch) in enumerate(app_dataloader):
-            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)
-
+            # Debug prints are already set for the whole session by set_model_debug_prints above
+            print(f"\n--- Training Batch {batch_idx+1}/{len(app_dataloader)} ---")
 
             src_batch, tgt_batch = src_batch.to(device_global), tgt_batch.to(device_global)
             decoder_input_tokens = src_batch[:, :-1] 
@@ -278,19 +268,20 @@ 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}"
+            # Print every batch to console due to full debug, but maybe less often to UI
+            print(log_line) 
             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"
         
-        # 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
     
-    set_model_debug_prints(swck_model_global, False, False, False) # Ensure off after all training
+    # Set debug prints OFF after the entire training session for subsequent operations (like generation)
+    # unless generation itself re-enables them.
+    print("--- App: Training Session Finished. Setting debug prints OFF by default. ---")
+    set_model_debug_prints(swck_model_global, False, False, False) 
     swck_model_global.eval() 
     
     try:
@@ -316,7 +307,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, enable_gen_debug: bool): # Add debug toggle
+def generate_text_for_app(prompt_str, max_len_gen, temperature_gen): # Removed debug toggle, always ON
     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."
@@ -324,17 +315,19 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen, enable_gen_d
     swck_model_global.eval() 
     swck_model_global.set_wiring_phase(False) 
     
-    # 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)
+    # FULL DEBUG ON for generation
+    print("\n--- App: Generating Text (Full Debug ON) ---")
+    set_model_debug_prints(swck_model_global, True, True, True)
     
-    print(f"App: Generating for prompt: '{prompt_str}', max_len: {max_len_gen}, temp: {temperature_gen}, Debug: {enable_gen_debug}")
+    print(f"App: Generating for prompt: '{prompt_str}', max_len: {max_len_gen}, temp: {temperature_gen}")
 
     tokens = [SOS_TOKEN] + [word_to_idx_global.get(w, UNK_TOKEN) for w in prompt_str.lower().split()]
     generated_ids_app = list(tokens)
-    debug_info_lines = [f"Prompt tokens: {generated_ids_app}"]
+    debug_info_lines = [f"Prompt tokens: {generated_ids_app}"] # For UI
 
     with torch.no_grad():
         for i in range(int(max_len_gen)): 
+            print(f"\n--- Generation Step {i+1} ---") # Console log for each step
             context_start_idx = max(0, len(generated_ids_app) - SEQ_LEN_APP)
             current_context_ids = generated_ids_app[context_start_idx:]
             
@@ -355,11 +348,14 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen, enable_gen_d
 
             if next_token_id == EOS_TOKEN:
                 debug_info_lines.append(f"Step {i+1}: EOS token encountered.")
+                print(f"Step {i+1}: EOS token encountered.")
                 break
             generated_ids_app.append(next_token_id)
             
-            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)
+            current_word = idx_to_word_global.get(next_token_id, UNK_TOKEN_STR)
+            print(f"  ==> Generated token {i+1}: '{current_word}' (ID: {next_token_id})") # Console log
+
+            if i < 10 : # UI debug info is still limited 
                 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:
                     b0_ent = entropy_report_infer['block_output_entropies'][0].item()
@@ -380,11 +376,12 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen, enable_gen_d
 
     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)
+    print("--- App: Generation Finished. Setting debug prints OFF by default. ---")
+    set_model_debug_prints(swck_model_global, False, False, False) # Turn off after this call
     return final_text, debug_output_str
 
-# Load model once on app startup. Set enable_initial_debug=False for cleaner startup logs.
+# Initialize model with debug OFF for initial startup to keep logs clean, 
+# will be turned ON by training/generation functions.
 initial_load_status = initialize_or_load_model_app(enable_initial_debug=False) 
 
 with gr.Blocks(title="SWCK Conceptual Demo") as demo:
@@ -392,7 +389,7 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
     
     gr.Markdown(f"""
     # Self-Wired Conscious Kernel (SWCK) - Conceptual Demo
-    This demo showcases a conceptual text generation model.
+    This demo showcases a conceptual text generation model with **FULL KERNEL DEBUGGING ON by default** for training and generation (output to Space console logs).
     Seed Phrase: "{SEED_PHRASE_APP[:100]}..." | Seed Number: "{SEED_NUMBER_STR_APP}".
     (Note: If checkpoint is not found or fails to load, an *untrained* model is used.)
     """)
@@ -401,7 +398,7 @@ 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)
+                # Removed debug checkbox as it's on by default for console
             with gr.Row():
                 generate_button = gr.Button("Generate", scale=1)
             with gr.Row():
@@ -412,21 +409,21 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
             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. Full Kernel Debug will be printed to console for the FIRST BATCH of the FIRST WIRING EPOCH ONLY.")
+            gr.Markdown("WARNING: In-app training is EXTREMELY slow. **Full Kernel Debug will be printed to console for ALL batches/epochs.** Model state persists only for this session unless saved manually.")
             with gr.Row():
-                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_epochs_slider = gr.Slider(minimum=1, maximum=2, value=1, step=1, label="Number of Training Epochs (1-2 for demo)") 
+                train_batch_size_slider = gr.Slider(minimum=1, maximum=2, value=1, step=1, label="Training Batch Size (1-2 for demo)") 
                 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 (summary):", lines=10, interactive=False,show_label=True )
+            start_training_button = gr.Button("Start Short Training Session (Full Debug to Console)")
+            training_status_output = gr.Textbox(label="Training Log / Status (summary to UI):", 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, enable_generation_debug_checkbox], # Added checkbox
+        fn=generate_text_for_app, # Removed enable_gen_debug from inputs
+        inputs=[prompt_input, max_len_slider, temp_slider], 
         outputs=[output_text, debug_text_area]
     )
     
@@ -438,6 +435,4 @@ 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) 
+    demo.launch(debug=True) # Gradio server debug