Update app.py

app.py

@@ -57,17 +57,6 @@ OVERALL_OUTPUT_ENTROPY_REG_WEIGHT_APP = 0.01
 GATE_SPARSITY_LOSS_WEIGHT_APP = 0.001
 WIRING_PHASE_EPOCHS_APP = 1 
 
-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
     print("App: Building vocabulary...")
@@ -84,8 +73,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
 
-# CORRECTED FUNCTION DEFINITION
-def initialize_or_load_model_app(enable_initial_debug=True): 
+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,19 +92,19 @@ def initialize_or_load_model_app(enable_initial_debug=True):
         'num_sub_modules_per_block': NUM_SUB_MODULES_PER_BLOCK_APP
     }
     
-    if enable_initial_debug:
-        print("App: Initializing SWCKModel with FULL DEBUG ON by default for init...")
-    
-    # Temporarily disable sub-component debug before SWCKModel init if enable_initial_debug is False,
-    # so SWCKModel's own init prints don't get mixed with sub-component init prints prematurely.
-    # SeedParser's internal debug_prints_enabled will control its own prints during its __init__.
+    print("App: Initializing SWCKModel. Debug prints are ON by default in model components.")
     
     swck_model_global = SWCKModel(**model_args).to(device_global)
-    # Now set the debug states for all components based on enable_initial_debug
-    set_model_debug_prints(swck_model_global, 
-                           seed_parser_debug=enable_initial_debug, 
-                           block_debug=enable_initial_debug, 
-                           model_debug=enable_initial_debug)
+    # Ensure debug flags are True on all components after initialization
+    # (assuming model.py might have them False by default, this makes them True)
+    if swck_model_global:
+        swck_model_global.debug_prints_enabled = True
+        if hasattr(swck_model_global, 'seed_parser'):
+            swck_model_global.seed_parser.debug_prints_enabled = True
+        if hasattr(swck_model_global, 'adaptive_blocks'):
+            for block in swck_model_global.adaptive_blocks:
+                block.debug_prints_enabled = True
+        print("App: Confirmed debug prints ON for SWCKModel and its components.")
 
 
     if os.path.exists(CHECKPOINT_FILENAME):
@@ -141,28 +129,30 @@ def initialize_or_load_model_app(enable_initial_debug=True):
             else:
                 print("App: word_to_idx not in checkpoint, using app's rebuilt vocab.")
             
-            # Ensure debug states are correctly set after loading
-            set_model_debug_prints(swck_model_global, 
-                                   seed_parser_debug=enable_initial_debug, 
-                                   block_debug=enable_initial_debug, 
-                                   model_debug=enable_initial_debug)
+            # After loading, ensure debug flags are still True
+            if swck_model_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 block in swck_model_global.adaptive_blocks: block.debug_prints_enabled = True
+                print("App: Re-confirmed debug prints ON after loading checkpoint.")
+
 
             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 with debug state: {enable_initial_debug}.")
+            print(f"App: Error loading model from checkpoint: {e}. Re-initializing new model.")
             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 swck_model_global: # Ensure debug is on for the new instance too
+                swck_model_global.debug_prints_enabled = True
+                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
             optimizer_global = optim.AdamW(swck_model_global.parameters(), lr=0.001)
-            model_load_status_global = f"Error loading checkpoint. Using new (untrained) model with debug: {enable_initial_debug}."
+            model_load_status_global = "Error loading checkpoint. Using new (untrained) model."
     else:
-        print(f"App: Checkpoint {CHECKPOINT_FILENAME} not found. Initializing new model with debug state: {enable_initial_debug}.")
-        # set_model_debug_prints was already called for a new model above
+        print(f"App: Checkpoint {CHECKPOINT_FILENAME} not found. Initializing new model.")
+        # Debug flags already set for a new model instance above
         optimizer_global = optim.AdamW(swck_model_global.parameters(), lr=0.001)
-        model_load_status_global = f"Initialized a new (untrained) model with debug: {enable_initial_debug}."
+        model_load_status_global = "Initialized a new (untrained) model."
     
     swck_model_global.eval() 
     return model_load_status_global
@@ -199,15 +189,14 @@ 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 (Full Debug ON for ALL batches/epochs) ---")
+    print("\n--- App: Starting Short Training Session (Full Debug ON for ALL batches/epochs by default) ---")
     progress(0, desc="Preparing training data...")
     
-    set_model_debug_prints(swck_model_global, True, True, True) # DEBUG ALWAYS ON FOR TRAINING
+    # Model debug flags are assumed to be already ON from initialize_or_load_model_app()
 
     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) 
         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)
@@ -226,10 +215,11 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
     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
-        print(f"\n>>> EPOCH {epoch+1} - Starting with Full Debug for all batches <<<")
+        # No need to toggle debug here; it's globally on for the model instance
 
         for batch_idx, (src_batch, tgt_batch) in enumerate(app_dataloader):
-            print(f"\n--- Training Batch {batch_idx+1}/{len(app_dataloader)} ---") # Explicit batch print
+            # Print statements within model.py's forward methods will now trigger automatically
+            print(f"\n--- Training Batch {batch_idx+1}/{len(app_dataloader)} (Epoch {epoch+1}) ---") 
 
             src_batch, tgt_batch = src_batch.to(device_global), tgt_batch.to(device_global)
             decoder_input_tokens = src_batch[:, :-1] 
@@ -277,17 +267,17 @@ 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(log_line) 
+            print(log_line) # This will go to console
             if batch_idx % max(1, len(app_dataloader)//2) == 0 or batch_idx == len(app_dataloader)-1 : 
-                training_log_output += log_line + "\n"
+                training_log_output += log_line + "\n" # Summary to UI
         
         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
     
-    print("--- App: Training Session Finished. Setting debug prints OFF by default. ---")
-    set_model_debug_prints(swck_model_global, False, False, False) 
+    print("--- App: Training Session Finished. Debug prints remain ON for the model instance. ---")
+    # No need to turn off debugs here if they are meant to be globally on for the app session
     swck_model_global.eval() 
     
     try:
@@ -321,9 +311,8 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
     swck_model_global.eval() 
     swck_model_global.set_wiring_phase(False) 
     
-    print("\n--- App: Generating Text (Full Debug ON) ---")
-    set_model_debug_prints(swck_model_global, True, True, True) # DEBUG ALWAYS ON FOR GENERATION
-    
+    # Model debug flags are assumed to be already ON globally from initialize_or_load_model_app()
+    print("\n--- App: Generating Text (Full Debug ON by default) ---")
     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()]
@@ -332,6 +321,7 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
 
     with torch.no_grad():
         for i in range(int(max_len_gen)): 
+            # Print statements inside SWCKModel's forward and AdaptiveBlock's forward will trigger
             print(f"\n--- Generation Step {i+1} ---") 
             context_start_idx = max(0, len(generated_ids_app) - SEQ_LEN_APP)
             current_context_ids = generated_ids_app[context_start_idx:]
@@ -381,11 +371,11 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
 
     debug_output_str = "\n".join(debug_info_lines)
     
-    print("--- App: Generation Finished. Setting debug prints OFF by default. ---")
-    set_model_debug_prints(swck_model_global, False, False, False) 
+    # Debug flags remain ON for the model instance for subsequent calls unless changed elsewhere
+    print("--- App: Generation Finished. Debug prints remain ON for the model instance. ---")
     return final_text, debug_output_str
 
-# Initialize model. Set enable_initial_debug=True for verbose init logs.
+# Initialize model with debug ON by default for the whole app session
 initial_load_status = initialize_or_load_model_app(enable_initial_debug=True) 
 
 with gr.Blocks(title="SWCK Conceptual Demo") as demo:
@@ -393,7 +383,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 with **FULL KERNEL DEBUGGING ON by default** for training and generation (output to Space console logs).
+    This demo showcases a conceptual text generation model with **FULL KERNEL DEBUGGING ON by default** for all operations (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.)
     """)
@@ -402,9 +392,8 @@ 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)
-                # Removed debug checkbox from here
             with gr.Row():
-                generate_button = gr.Button("Generate (Full Debug to Console)", scale=1) # Updated button label
+                generate_button = gr.Button("Generate (Full Debug to Console)", 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)")
@@ -427,7 +416,7 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
 
     generate_button.click(
         fn=generate_text_for_app,
-        inputs=[prompt_input, max_len_slider, temp_slider], # Removed checkbox from inputs
+        inputs=[prompt_input, max_len_slider, temp_slider], 
         outputs=[output_text, debug_text_area]
     )