Update app.py

app.py

@@ -57,6 +57,17 @@ 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_component in model.adaptive_blocks: # Renamed to avoid conflict
+                block_component.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...")
@@ -73,7 +84,8 @@ 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(): 
+# CORRECTED FUNCTION DEFINITION: Added enable_initial_debug parameter
+def initialize_or_load_model_app(enable_initial_debug=True): 
     global swck_model_global, optimizer_global, word_to_idx_global, idx_to_word_global, \
            VOCAB_SIZE_APP, model_load_status_global
 
@@ -92,19 +104,14 @@ def initialize_or_load_model_app():
         'num_sub_modules_per_block': NUM_SUB_MODULES_PER_BLOCK_APP
     }
     
-    print("App: Initializing SWCKModel. Debug prints are ON by default in model components.")
+    if enable_initial_debug: # This print will now work correctly
+        print("App: Initializing SWCKModel with FULL DEBUG ON by default for init...")
     
     swck_model_global = SWCKModel(**model_args).to(device_global)
-    # 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.")
+    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):
@@ -129,30 +136,26 @@ def initialize_or_load_model_app():
             else:
                 print("App: word_to_idx not in checkpoint, using app's rebuilt vocab.")
             
-            # 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.")
-
+            set_model_debug_prints(swck_model_global, 
+                                   seed_parser_debug=enable_initial_debug, 
+                                   block_debug=enable_initial_debug, 
+                                   model_debug=enable_initial_debug)
 
             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.")
             swck_model_global = SWCKModel(**model_args).to(device_global)
-            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
+            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."
+            model_load_status_global = f"Error loading checkpoint. Using new (untrained) model with debug: {enable_initial_debug}."
     else:
-        print(f"App: Checkpoint {CHECKPOINT_FILENAME} not found. Initializing new model.")
-        # Debug flags already set for a new model instance above
+        print(f"App: Checkpoint {CHECKPOINT_FILENAME} not found. Initializing new model with debug state: {enable_initial_debug}.")
         optimizer_global = optim.AdamW(swck_model_global.parameters(), lr=0.001)
-        model_load_status_global = "Initialized a new (untrained) model."
+        model_load_status_global = f"Initialized a new (untrained) model with debug: {enable_initial_debug}."
     
     swck_model_global.eval() 
     return model_load_status_global
@@ -192,11 +195,13 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
     print("\n--- App: Starting Short Training Session (Full Debug ON for ALL batches/epochs by default) ---")
     progress(0, desc="Preparing training data...")
     
-    # Model debug flags are assumed to be already ON from initialize_or_load_model_app()
+    # 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 before training starts
         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)
@@ -215,10 +220,9 @@ 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
-        # No need to toggle debug here; it's globally on for the model instance
+        print(f"\n>>> EPOCH {epoch+1} - Starting with Full Debug for all batches <<<")
 
         for batch_idx, (src_batch, tgt_batch) in enumerate(app_dataloader):
-            # 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)
@@ -267,17 +271,18 @@ 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) # This will go to console
+            print(log_line) 
             if batch_idx % max(1, len(app_dataloader)//2) == 0 or batch_idx == len(app_dataloader)-1 : 
-                training_log_output += log_line + "\n" # Summary to UI
+                training_log_output += log_line + "\n" 
         
         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
     
+    # After training, leave debug ON as per request for "default ON" for the app instance.
+    # If you wanted it off after training, you'd call set_model_debug_prints(..., 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:
@@ -311,7 +316,7 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
     swck_model_global.eval() 
     swck_model_global.set_wiring_phase(False) 
     
-    # Model debug flags are assumed to be already ON globally from initialize_or_load_model_app()
+    # Debug is assumed to be ON from initialization for the model instance
     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}")
 
@@ -321,7 +326,6 @@ 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:]
@@ -371,11 +375,11 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
 
     debug_output_str = "\n".join(debug_info_lines)
     
-    # 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. ---")
+    # No need to turn off debugs if they are globally ON for the app session
     return final_text, debug_output_str
 
-# Initialize model with debug ON by default for the whole app session
+# Initialize model with debug ON by default for the entire app session
 initial_load_status = initialize_or_load_model_app(enable_initial_debug=True) 
 
 with gr.Blocks(title="SWCK Conceptual Demo") as demo: