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_fulldebug.pth.tar" # New checkpoint name
+CHECKPOINT_FILENAME = "swck_model_conceptual_app_fulldebug.pth.tar" 
 
 MAIN_LOSS_WEIGHT_APP = 1.0
 BLOCK_TARGET_ENTROPY_LOSS_WEIGHT_APP = 0.02
@@ -84,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
+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
 
@@ -103,17 +104,19 @@ def initialize_or_load_model_app():
         'num_sub_modules_per_block': NUM_SUB_MODULES_PER_BLOCK_APP
     }
     
-    print("App: Initializing SWCKModel with FULL DEBUG ON by default for init...")
+    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__.
+    
     swck_model_global = SWCKModel(**model_args).to(device_global)
-    # 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.")
+    # 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)
 
 
     if os.path.exists(CHECKPOINT_FILENAME):
@@ -137,21 +140,29 @@ def initialize_or_load_model_app():
                     print("App: Checkpoint vocab seems invalid, using app's rebuilt vocab.")
             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)
 
             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 ON.")
+            print(f"App: Error loading model from checkpoint: {e}. Re-initializing new model with debug state: {enable_initial_debug}.")
             swck_model_global = SWCKModel(**model_args).to(device_global)
-            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
+            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 with debug ON."
+            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 with debug ON.")
+        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
         optimizer_global = optim.AdamW(swck_model_global.parameters(), lr=0.001)
-        model_load_status_global = "Initialized a new (untrained) model with debug ON."
+        model_load_status_global = f"Initialized a new (untrained) model with debug: {enable_initial_debug}."
     
     swck_model_global.eval() 
     return model_load_status_global
@@ -191,13 +202,12 @@ 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) ---")
     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)
+    set_model_debug_prints(swck_model_global, True, True, True) # DEBUG ALWAYS ON FOR TRAINING
 
     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
+        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)
@@ -219,8 +229,7 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
         print(f"\n>>> EPOCH {epoch+1} - Starting with Full Debug for all batches <<<")
 
         for batch_idx, (src_batch, tgt_batch) in enumerate(app_dataloader):
-            # 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)} ---")
+            print(f"\n--- Training Batch {batch_idx+1}/{len(app_dataloader)} ---") # Explicit batch print
 
             src_batch, tgt_batch = src_batch.to(device_global), tgt_batch.to(device_global)
             decoder_input_tokens = src_batch[:, :-1] 
@@ -268,7 +277,6 @@ 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 : 
                 training_log_output += log_line + "\n"
@@ -278,8 +286,6 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
         print(epoch_summary)
         training_log_output += epoch_summary
     
-    # 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() 
@@ -307,7 +313,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): # Removed debug toggle, always ON
+def generate_text_for_app(prompt_str, max_len_gen, temperature_gen): 
     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."
@@ -315,19 +321,18 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen): # Removed d
     swck_model_global.eval() 
     swck_model_global.set_wiring_phase(False) 
     
-    # FULL DEBUG ON for generation
     print("\n--- App: Generating Text (Full Debug ON) ---")
-    set_model_debug_prints(swck_model_global, True, True, True)
+    set_model_debug_prints(swck_model_global, True, True, True) # DEBUG ALWAYS ON FOR GENERATION
     
     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}"] # For UI
+    debug_info_lines = [f"Prompt tokens: {generated_ids_app}"] 
 
     with torch.no_grad():
         for i in range(int(max_len_gen)): 
-            print(f"\n--- Generation Step {i+1} ---") # Console log for each step
+            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:]
             
@@ -353,9 +358,9 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen): # Removed d
             generated_ids_app.append(next_token_id)
             
             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
+            print(f"  ==> Generated token {i+1}: '{current_word}' (ID: {next_token_id})") 
 
-            if i < 10 : # UI debug info is still limited 
+            if i < 10 : 
                 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()
@@ -377,12 +382,11 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen): # Removed d
     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) # Turn off after this call
+    set_model_debug_prints(swck_model_global, False, False, False) 
     return final_text, debug_output_str
 
-# 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() 
+# Initialize model. Set enable_initial_debug=True for verbose init logs.
+initial_load_status = initialize_or_load_model_app(enable_initial_debug=True) 
 
 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")
@@ -398,9 +402,9 @@ 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 as it's on by default for console
+                # Removed debug checkbox from here
             with gr.Row():
-                generate_button = gr.Button("Generate", scale=1)
+                generate_button = gr.Button("Generate (Full Debug to Console)", scale=1) # Updated button label
             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)")
@@ -422,8 +426,8 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
         return f"**Model Status:** {model_load_status_global}"
 
     generate_button.click(
-        fn=generate_text_for_app, # Removed enable_gen_debug from inputs
-        inputs=[prompt_input, max_len_slider, temp_slider], 
+        fn=generate_text_for_app,
+        inputs=[prompt_input, max_len_slider, temp_slider], # Removed checkbox from inputs
         outputs=[output_text, debug_text_area]
     )
     
@@ -435,4 +439,4 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
     
 
 if __name__ == "__main__":
-    demo.launch(debug=True) # Gradio server debug
+    demo.launch(debug=True)