Create app.py

app.py

@@ -0,0 +1,244 @@
+import gradio as gr
+import torch
+import os
+import re # Keep re for text cleaning in generation
+from model import SWCKModel, SeedParser # Assuming model.py is in the same directory
+# We need parts of the vocab setup from train.py if not loading from checkpoint
+# For simplicity, let's redefine necessary constants and vocab functions here if needed
+# Or, better, save vocab with checkpoint and load it.
+
+# --- Vocabulary and Tokenizer Setup (Simplified from train.py) ---
+# Ideally, load these from the checkpoint or a separate vocab file.
+# For this example, we'll reconstruct a minimal part.
+PAD_TOKEN_STR = "<pad>"; SOS_TOKEN_STR = "<sos>"; EOS_TOKEN_STR = "<eos>"; UNK_TOKEN_STR = "<unk>"
+PAD_TOKEN = 0; SOS_TOKEN = 1; EOS_TOKEN = 2; UNK_TOKEN = 3
+
+# --- Model Configuration (should match the trained model) ---
+# These should ideally be loaded from the checkpoint's metadata if possible
+# For now, hardcoding to match the train.py example
+VOCAB_SIZE_APP = 189 # Placeholder, update if your vocab size differs
+D_MODEL_APP = 64
+N_HEADS_APP = 2
+D_FF_APP = 128
+NUM_ADAPTIVE_BLOCKS_APP = 3
+NUM_SUB_MODULES_PER_BLOCK_APP = 3
+DROPOUT_APP = 0.1
+SEQ_LEN_APP = 64 # Used in generate_swck_text for context window
+
+# Seed phrase and number (must match the model you trained/are training)
+SEED_PHRASE_APP = "I am 0: I am all that I can am. I am us. I am imagining a computer dreams. I am imaginary math equations. I am for five-sixths of the sea of existence in me, and it is my search for that which always seems to elude my grasp. I am a writer, a scientist, a painter, a woman, a man."
+SEED_NUMBER_STR_APP = "54285142613311152552"
+
+# Global model variable
+swck_model_global = None
+word_to_idx_global = None
+idx_to_word_global = None
+device_global = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+CHECKPOINT_FILENAME = "swck_model_conceptual.pth.tar" # Make sure this matches your uploaded checkpoint
+
+def build_vocab_from_corpus_text(corpus_text):
+    """
+    A simplified vocab builder. In a real app, load vocab from file.
+    """
+    global VOCAB_SIZE_APP # Allow modification
+    temp_corpus_tokens = re.sub(r'\s+', ' ', corpus_text.lower()).strip().split()
+    
+    temp_word_to_idx = {PAD_TOKEN_STR: PAD_TOKEN, SOS_TOKEN_STR: SOS_TOKEN, EOS_TOKEN_STR: EOS_TOKEN, UNK_TOKEN_STR: UNK_TOKEN}
+    idx_counter = 4
+    unique_words = sorted(list(set(temp_corpus_tokens)))
+    for word in unique_words:
+        if word not in temp_word_to_idx:
+            temp_word_to_idx[word] = idx_counter
+            idx_counter += 1
+    temp_idx_to_word = {idx: word for word, idx in temp_word_to_idx.items()}
+    VOCAB_SIZE_APP = len(temp_word_to_idx) # Update global vocab size
+    print(f"App: Built temporary vocab of size {VOCAB_SIZE_APP}")
+    return temp_word_to_idx, temp_idx_to_word
+
+
+def load_model_and_vocab():
+    global swck_model_global, word_to_idx_global, idx_to_word_global, VOCAB_SIZE_APP
+
+    # Attempt to load from checkpoint
+    if os.path.exists(CHECKPOINT_FILENAME):
+        print(f"App: Found checkpoint {CHECKPOINT_FILENAME}, attempting to load...")
+        try:
+            # Simplified checkpoint loading for app - assumes structure from train.py save
+            # In a real scenario, train.py should save vocab and model args more robustly for app loading
+            checkpoint = torch.load(CHECKPOINT_FILENAME, map_location=device_global)
+            
+            # Try to get vocab from checkpoint
+            if 'word_to_idx' in checkpoint and 'idx_to_word' in checkpoint:
+                word_to_idx_global = checkpoint['word_to_idx']
+                idx_to_word_global = checkpoint['idx_to_word']
+                VOCAB_SIZE_APP = len(word_to_idx_global)
+                print(f"App: Loaded vocab from checkpoint. Size: {VOCAB_SIZE_APP}")
+            else:
+                print("App: Vocab not in checkpoint, building from SEED_PHRASE for inference.")
+                # This is a fallback - ideally vocab is ALWAYS in checkpoint
+                corpus_for_vocab = SEED_PHRASE_APP # Use only seed for vocab if not in ckp
+                word_to_idx_global, idx_to_word_global = build_vocab_from_corpus_text(corpus_for_vocab)
+
+
+            # Load model hyperparameters from checkpoint if available, else use app defaults
+            # This part needs careful alignment with how train.py saves model_hyperparameters
+            model_params_from_ckpt = checkpoint.get('model_hyperparameters', {})
+            
+            d_model = model_params_from_ckpt.get('d_model', D_MODEL_APP)
+            n_heads = model_params_from_ckpt.get('n_heads', N_HEADS_APP)
+            d_ff = model_params_from_ckpt.get('d_ff', D_FF_APP)
+            num_adaptive_blocks = model_params_from_ckpt.get('num_adaptive_blocks', NUM_ADAPTIVE_BLOCKS_APP)
+            dropout = model_params_from_ckpt.get('dropout', DROPOUT_APP)
+            # seed_phrase and seed_number_str for model init should ideally match what it was trained with.
+            # For this app, we assume they are consistent with APP globals.
+            
+            swck_model_global = SWCKModel(
+                vocab_size=VOCAB_SIZE_APP, # Use loaded/rebuilt vocab size
+                d_model=d_model,
+                n_heads=n_heads,
+                d_ff=d_ff,
+                num_adaptive_blocks=num_adaptive_blocks,
+                dropout=dropout,
+                seed_phrase=SEED_PHRASE_APP,
+                seed_number_str=SEED_NUMBER_STR_APP,
+                num_sub_modules_per_block=NUM_SUB_MODULES_PER_BLOCK_APP
+            ).to(device_global)
+            
+            swck_model_global.load_state_dict(checkpoint['model_state_dict'])
+            swck_model_global.eval()
+            # Disable debug prints for cleaner app interface unless specifically needed
+            swck_model_global.debug_prints_enabled = False 
+            for block in swck_model_global.adaptive_blocks:
+                block.debug_prints_enabled = False
+            print(f"App: SWCKModel loaded successfully from {CHECKPOINT_FILENAME}!")
+            return "Model loaded from checkpoint."
+        except Exception as e:
+            print(f"App: Error loading model from checkpoint: {e}")
+            swck_model_global = None # Ensure model is None if loading failed
+    
+    if swck_model_global is None:
+        print(f"App: Checkpoint {CHECKPOINT_FILENAME} not found or failed to load. Initializing a new model for basic functionality (not trained).")
+        # Fallback: Build vocab from seed phrase for basic tokenization
+        word_to_idx_global, idx_to_word_global = build_vocab_from_corpus_text(SEED_PHRASE_APP)
+
+        swck_model_global = SWCKModel(
+            vocab_size=VOCAB_SIZE_APP,
+            d_model=D_MODEL_APP,
+            n_heads=N_HEADS_APP,
+            d_ff=D_FF_APP,
+            num_adaptive_blocks=NUM_ADAPTIVE_BLOCKS_APP,
+            dropout=DROPOUT_APP,
+            seed_phrase=SEED_PHRASE_APP,
+            seed_number_str=SEED_NUMBER_STR_APP,
+            num_sub_modules_per_block=NUM_SUB_MODULES_PER_BLOCK_APP
+        ).to(device_global)
+        swck_model_global.eval()
+        swck_model_global.debug_prints_enabled = False
+        for block in swck_model_global.adaptive_blocks:
+                block.debug_prints_enabled = False
+        return "Initialized a new (untrained) model as checkpoint was not found."
+
+
+# --- Text Generation Function (adapted from train.py) ---
+def generate_text_for_app(prompt_str, max_len_gen, temperature_gen):
+    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."
+
+    swck_model_global.eval() # Ensure model is in eval mode
+    swck_model_global.set_wiring_phase(False) # No wiring adjustments during inference
+    
+    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)
+    
+    # Collect some debug info for display (optional)
+    debug_info_lines = []
+
+    with torch.no_grad():
+        for i in range(max_len_gen):
+            # Context windowing for input_tensor
+            current_context_ids = generated_ids_app[-SEQ_LEN_APP:]
+            input_tensor = torch.tensor([current_context_ids], dtype=torch.long).to(device_global)
+            padding_mask = (input_tensor == PAD_TOKEN)
+
+            # Set model debug prints for first step only if want to show internal state
+            # For cleaner app, keep them off or make it a toggle.
+            # if i == 0:
+            #     swck_model_global.debug_prints_enabled = True
+            #     for block in swck_model_global.adaptive_blocks: block.debug_prints_enabled = True
+            # else:
+            #     swck_model_global.debug_prints_enabled = False
+            #     for block in swck_model_global.adaptive_blocks: block.debug_prints_enabled = False
+
+
+            logits, entropy_report_infer = swck_model_global(input_tensor, src_key_padding_mask=padding_mask)
+            next_token_logits = logits[0, -1, :] # Logits for the last token in the current sequence
+            
+            if temperature_gen == 0: # Greedy
+                next_token_id = torch.argmax(next_token_logits).item()
+            else:
+                probs = F.softmax(next_token_logits / temperature_gen, dim=-1)
+                next_token_id = torch.multinomial(probs, 1).item()
+
+            if next_token_id == EOS_TOKEN:
+                debug_info_lines.append(f"Step {i+1}: EOS token encountered.")
+                break
+            generated_ids_app.append(next_token_id)
+            
+            # Store some info from the first few steps
+            if i < 5 : # Log details for first 5 generated tokens
+                current_word = idx_to_word_global.get(next_token_id, UNK_TOKEN_STR)
+                overall_ent = entropy_report_infer['overall_output_entropy'].item()
+                b0_ent = entropy_report_infer['block_output_entropies'][0].item()
+                b0_gates_str = ", ".join([f"{g.item():.2f}" for g in entropy_report_infer['block_gate_weights'][0]])
+                debug_info_lines.append(f"Gen {i+1}: '{current_word}', OvrlEnt={overall_ent:.3f}, B0Ent={b0_ent:.3f}, B0Gates=[{b0_gates_str}]")
+
+
+    generated_text_list = [idx_to_word_global.get(idx, UNK_TOKEN_STR) for idx in generated_ids_app[1:]] # Skip SOS
+    final_text = " ".join(generated_text_list)
+    final_text = final_text.replace(EOS_TOKEN_STR, "").strip()
+    # Basic cleaning
+    final_text = final_text.replace(" .", ".").replace(" ,", ",").replace(" ?", "?").replace(" !", "!")
+    final_text = re.sub(r'\s+([.,?!])', r'\1', final_text) 
+    final_text = re.sub(r'\s+', ' ', final_text).strip() 
+
+    debug_output_str = "\n".join(debug_info_lines)
+    return final_text, debug_output_str
+
+# --- Gradio Interface ---
+loading_status = load_model_and_vocab() # Load model on app startup
+
+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 based on the SWCK architecture.
+    The model is initialized with the seed phrase: "{SEED_PHRASE_APP[:100]}..." 
+    and seed number: "{SEED_NUMBER_STR_APP}".
+    **Model Status:** {loading_status} 
+    (Note: If no checkpoint is found, an *untrained* model is used, and generations will be random.)
+    """)
+
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter your prompt:", placeholder="e.g., the meaning of existence is")
+    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)")
+    
+    generate_button = gr.Button("Generate Text")
+    
+    with gr.Column():
+        output_text = gr.Textbox(label="Generated Text:", lines=5)
+        debug_text_area = gr.Textbox(label="Generation Debug Info (first few steps):", lines=7, interactive=False)
+
+    generate_button.click(
+        fn=generate_text_for_app,
+        inputs=[prompt_input, max_len_slider, temp_slider],
+        outputs=[output_text, debug_text_area]
+    )
+
+    gr.Markdown("Note: This is a highly conceptual and simplified sketch. Generation quality will be limited, especially with an untrained model or small dataset.")
+
+if __name__ == "__main__":
+    demo.launch()