Update app.py

app.py

@@ -93,10 +93,10 @@ def initialize_or_load_model_app():
     }
 
     swck_model_global = SWCKModel(**model_args).to(device_global)
-    swck_model_global.debug_prints_enabled = True # Top-level model debug
+    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 # Block-level debug
+        block.debug_prints_enabled = True 
         # print(f"App: Debug prints explicitly enabled for AdaptiveBlock {i}")
 
 
@@ -112,18 +112,16 @@ def initialize_or_load_model_app():
 
             if 'word_to_idx' in checkpoint:
                 loaded_w2i = checkpoint['word_to_idx']
-                # Basic check, could be more robust
                 if isinstance(loaded_w2i, dict) and len(loaded_w2i) > 4: 
                     word_to_idx_global = loaded_w2i
                     idx_to_word_global = {v: k for k,v in loaded_w2i.items()}
-                    VOCAB_SIZE_APP = len(word_to_idx_global) # Ensure vocab size reflects loaded
+                    VOCAB_SIZE_APP = len(word_to_idx_global) 
                     print(f"App: Overwrote vocab with checkpoint's vocab. New size: {VOCAB_SIZE_APP}")
                 else:
                     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.")
 
-
             model_load_status_global = f"Model loaded successfully from {CHECKPOINT_FILENAME}."
             print(model_load_status_global)
         except Exception as e:
@@ -148,11 +146,9 @@ class AppSWCKDataset(Dataset):
         self.seq_len = seq_len
         self.sos_id, self.eos_id, self.pad_id = sos_id, eos_id, pad_id
         self.samples = []
-        # Create overlapping sequences for language modeling
-        # Ensure target is seq_len for consistency with input to model.
-        for i in range(len(token_ids) - seq_len -1): # -1 to ensure target has full seq_len
-            input_seq = [self.sos_id] + token_ids[i : i + seq_len] # length seq_len + 1
-            target_seq = token_ids[i + 1 : i + seq_len + 1] + [self.eos_id] # length seq_len + 1
+        for i in range(len(token_ids) - seq_len -1): 
+            input_seq = [self.sos_id] + token_ids[i : i + seq_len] 
+            target_seq = token_ids[i + 1 : i + seq_len + 1] + [self.eos_id] 
             self.samples.append((input_seq, target_seq))
         print(f"AppSWCKDataset: Created {len(self.samples)} training samples for in-app training.")
 
@@ -198,39 +194,35 @@ 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
         
-        # Enable debug for first batch of first epoch
         first_batch_debug = (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) : # Disable after first batch for speed
+            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
 
 
             src_batch, tgt_batch = src_batch.to(device_global), tgt_batch.to(device_global)
-            decoder_input_tokens = src_batch[:, :-1] # Remove EOS from input
-            gold_standard_for_loss = tgt_batch[:, 1:] # Remove SOS from target
+            decoder_input_tokens = src_batch[:, :-1] 
+            gold_standard_for_loss = tgt_batch[:, 1:] 
             
             src_key_padding_mask = (decoder_input_tokens == PAD_TOKEN)
 
             optimizer_global.zero_grad()
             logits, entropy_report = swck_model_global(decoder_input_tokens, src_key_padding_mask=src_key_padding_mask)
             
-            # Ensure logits and gold_standard_for_loss are aligned for CrossEntropyLoss
-            # Logits: (B, S_len_in, VocabSize)
-            # Gold: (B, S_len_target)
-            # If S_len_in == S_len_target, it's fine.
             if logits.size(1) != gold_standard_for_loss.size(1):
-                # This can happen if seq len handling differs slightly, adjust shorter one
                 min_len = min(logits.size(1), gold_standard_for_loss.size(1))
-                logits_for_loss = logits[:, :min_len, :].contiguous()
+                logits_for_loss = logits[:, :min_len, :].contiguous() # ADDED .contiguous()
                 gold_for_loss_aligned = gold_standard_for_loss[:, :min_len].contiguous()
             else:
-                logits_for_loss = logits
-                gold_for_loss_aligned = gold_standard_for_loss
+                logits_for_loss = logits.contiguous() # ADDED .contiguous()
+                gold_for_loss_aligned = gold_standard_for_loss.contiguous()
 
             main_loss = criterion_main_app(logits_for_loss.view(-1, logits_for_loss.size(-1)), gold_for_loss_aligned.view(-1))
 
@@ -239,7 +231,7 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
                 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))
-                if entropy_report["block_output_entropies"]: # Avoid division by zero
+                if entropy_report["block_output_entropies"]: 
                     block_entropy_loss = block_entropy_loss / len(entropy_report["block_output_entropies"])
 
             overall_entropy_loss = entropy_report["overall_output_entropy"]
@@ -247,7 +239,7 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
             if entropy_report["block_gate_weights"]:
                 for gates_softmax_tensor in entropy_report["block_gate_weights"]:
                     gate_sparsity_loss += torch.mean(gates_softmax_tensor * torch.log(gates_softmax_tensor + 1e-9))
-                if entropy_report["block_gate_weights"]: # Avoid division by zero
+                if entropy_report["block_gate_weights"]: 
                      gate_sparsity_loss = - (gate_sparsity_loss / len(entropy_report["block_gate_weights"]))
 
             combined_loss = (MAIN_LOSS_WEIGHT_APP * main_loss +
@@ -261,12 +253,12 @@ 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)//2) == 0 or batch_idx == len(app_dataloader)-1 : # Log less frequently to UI
+            if batch_idx % max(1, len(app_dataloader)//5) == 0 or batch_idx == len(app_dataloader)-1 : 
                 print(log_line) 
                 training_log_output += log_line + "\n"
         
-        # Disable debug prints after the very first batch of the first epoch
         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
 
 
@@ -275,8 +267,8 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
         print(epoch_summary)
         training_log_output += epoch_summary
     
-    # Ensure debug prints are off after training session
     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
     swck_model_global.eval() 
     
@@ -310,8 +302,10 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
 
     swck_model_global.eval() 
     swck_model_global.set_wiring_phase(False) 
-    # Temporarily enable debug for generation if needed, then disable
-    # swck_model_global.debug_prints_enabled = True # For generation debug
+    
+    # Temporarily re-enable debug for generation if you want to inspect Space logs
+    # 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
     
     print(f"App: Generating for prompt: '{prompt_str}', max_len: {max_len_gen}, temp: {temperature_gen}")
@@ -321,9 +315,7 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
     debug_info_lines = [f"Prompt tokens: {generated_ids_app}"]
 
     with torch.no_grad():
-        for i in range(int(max_len_gen)): # Ensure max_len_gen is int
-            # Context windowing for input_tensor
-            # Take up to SEQ_LEN_APP tokens from the end of generated_ids_app
+        for i in range(int(max_len_gen)): 
             context_start_idx = max(0, len(generated_ids_app) - SEQ_LEN_APP)
             current_context_ids = generated_ids_app[context_start_idx:]
             
@@ -360,7 +352,6 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
                 else:
                     debug_info_lines.append(f"Gen {i+1}: '{current_word}', OvrlEnt={overall_ent:.3f}, No block entropy/gate report.")
 
-
     generated_text_list = [idx_to_word_global.get(idx, UNK_TOKEN_STR) for idx in generated_ids_app[1:]] 
     final_text = " ".join(generated_text_list)
     final_text = final_text.replace(EOS_TOKEN_STR, "").strip()
@@ -370,20 +361,21 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
 
     debug_output_str = "\n".join(debug_info_lines)
     
-    # Disable debug prints after generation
-    # swck_model_global.debug_prints_enabled = False
+    # swck_model_global.debug_prints_enabled = False # Disable after generation
+    # 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
     return final_text, debug_output_str
 
-# --- Gradio Interface ---
-initial_load_status = initialize_or_load_model_app() # Load model on app startup
+initial_load_status = initialize_or_load_model_app()
 
 with gr.Blocks(title="SWCK Conceptual Demo") as demo:
+    # Using a unique elem_id for the status Markdown
+    model_status_md = gr.Markdown(value=f"**Model Status:** {initial_load_status}", elem_id="model_status_md_123")
+    
     gr.Markdown(f"""
     # Self-Wired Conscious Kernel (SWCK) - Conceptual Demo
     This demo showcases a conceptual text generation model.
     Seed Phrase: "{SEED_PHRASE_APP[:100]}..." | Seed Number: "{SEED_NUMBER_STR_APP}".
-    **Model Status:** <span id="model_status_display">{initial_load_status}</span>
     (Note: If checkpoint is not found or fails to load, an *untrained* model is used.)
     """)
     
@@ -404,16 +396,13 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
             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")
-                # REMOVED format="%.1e"
                 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 )
 
-
-    model_status_md = gr.Markdown(value=f"**Model Status:** {model_load_status_global}")
-
-    def update_status_text(): # Helper to refresh status after training
+    def update_status_text_for_ui(): 
+        # This function will be called by .then() to get the new status string
         return f"**Model Status:** {model_load_status_global}"
 
     generate_button.click(
@@ -426,11 +415,8 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
         fn=run_short_training_session,
         inputs=[train_epochs_slider, train_batch_size_slider, train_lr_slider],
         outputs=[training_status_output]
-    ).then(fn=update_status_text, inputs=None, outputs=model_status_md)
+    ).then(fn=update_status_text_for_ui, inputs=None, outputs=model_status_md) # Update the Markdown component
     
 
 if __name__ == "__main__":
-    # The Gradio app launch options (like debug=True) are for local execution.
-    # On Hugging Face Spaces, these are typically controlled by the environment.
-    # The `print()` statements will go to the Space's console logs.
     demo.launch(debug=True)