Update app.py

app.py

@@ -2,11 +2,12 @@ import gradio as gr
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.utils.data import Dataset, DataLoader # For dummy training
+from torch.utils.data import Dataset, DataLoader 
 import os
 import re
-import time # For basic progress update
-from model import SWCKModel, SeedParser, EntropyEstimator # Assuming model.py is in the same directory
+import time 
+import torch.nn.functional as F # <<<<<<<<<<<< ADDED THIS IMPORT
+from model import SWCKModel, SeedParser, EntropyEstimator 
 
 # --- Vocabulary and Tokenizer Setup ---
 PAD_TOKEN_STR = "<pad>"; SOS_TOKEN_STR = "<sos>"; EOS_TOKEN_STR = "<eos>"; UNK_TOKEN_STR = "<unk>"
@@ -218,10 +219,10 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
             
             if logits.size(1) != gold_standard_for_loss.size(1):
                 min_len = min(logits.size(1), gold_standard_for_loss.size(1))
-                logits_for_loss = logits[:, :min_len, :].contiguous() # ADDED .contiguous()
+                logits_for_loss = logits[:, :min_len, :].contiguous() 
                 gold_for_loss_aligned = gold_standard_for_loss[:, :min_len].contiguous()
             else:
-                logits_for_loss = logits.contiguous() # ADDED .contiguous()
+                logits_for_loss = logits.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))
@@ -230,7 +231,7 @@ def run_short_training_session(num_epochs_app, batch_size_app, learning_rate_app
             if entropy_report["block_output_entropies"]:
                 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))
+                    block_entropy_loss += F.mse_loss(block_entropy_tensor, torch.tensor(target_entropy_val, device=device_global)) # Used F here
                 if entropy_report["block_output_entropies"]: 
                     block_entropy_loss = block_entropy_loss / len(entropy_report["block_output_entropies"])
 
@@ -303,11 +304,6 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
     swck_model_global.eval() 
     swck_model_global.set_wiring_phase(False) 
     
-    # 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}")
 
     tokens = [SOS_TOKEN] + [word_to_idx_global.get(w, UNK_TOKEN) for w in prompt_str.lower().split()]
@@ -328,7 +324,7 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
             if temperature_gen == 0: 
                 next_token_id = torch.argmax(next_token_logits).item()
             else:
-                probs = F.softmax(next_token_logits / temperature_gen, dim=-1)
+                probs = F.softmax(next_token_logits / temperature_gen, dim=-1) # Used F here
                 if probs.isnan().any() or probs.isinf().any() or torch.sum(probs).item() < 1e-9 : 
                     print(f"Warning: Invalid probabilities at step {i}. Using uniform.")
                     probs = torch.ones_like(next_token_logits) / next_token_logits.size(-1) 
@@ -361,15 +357,11 @@ def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
 
     debug_output_str = "\n".join(debug_info_lines)
     
-    # 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
 
 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"""
@@ -402,7 +394,6 @@ with gr.Blocks(title="SWCK Conceptual Demo") as demo:
             training_status_output = gr.Textbox(label="Training Log / Status:", lines=10, interactive=False,show_label=True )
 
     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(
@@ -415,7 +406,7 @@ 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_for_ui, inputs=None, outputs=model_status_md) # Update the Markdown component
+    ).then(fn=update_status_text_for_ui, inputs=None, outputs=model_status_md) 
     
 
 if __name__ == "__main__":