Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py (Hardened and Debuggable Version)
+# app.py (Final Version with gr.State for Robust State Management)
 import gradio as gr
 import tensorflow as tf
 import pickle
@@ -22,24 +22,27 @@ except Exception as e:
     print(f"FATAL ERROR loading files: {e}")
     successor_model, predecessor_model = None, None
 
-# --- 2. THE CORE PREDICTION LOGIC (MODIFIED) ---
-# This function now receives the actual model and tokenizer objects
-def predict_next_state(model, tokenizers, current_unit, current_analogy, current_commentary):
-    if not model or not tokenizers:
-        return "Error: A required model or tokenizer is not loaded.", "Check server logs.", "---"
+# --- 2. THE CORE PREDICTION LOGIC ---
+# This function is the same, but it will now receive its input from the reliable gr.State
+def predict_next_state(model, tokenizers, current_state_dict):
+    if not model or not tokenizers or not current_state_dict:
+        return {"error": "Model or state is not loaded"}, "Error", "Error", "Error"
 
-    # Prepare input data
-    input_data = {'current_unit_name': [current_unit], 'current_analogy': [current_analogy], 'current_commentary': [current_commentary]}
+    # Prepare input data from the state dictionary
+    input_data = {
+        'current_unit_name': [current_state_dict['unit_name']],
+        'current_analogy': [current_state_dict['analogy']],
+        'current_commentary': [current_state_dict['commentary']]
+    }
     processed_input = {}
     for col, text_list in input_data.items():
         sequences = tokenizers[col].texts_to_sequences(text_list)
         padded_sequences = tf.keras.preprocessing.sequence.pad_sequences(sequences, maxlen=MAX_SEQ_LENGTH, padding='post')
         processed_input[col] = padded_sequences
         
-    # Get AI prediction
     predictions = model.predict(processed_input)
     
-    # Decode prediction back to text
+    # Decode prediction
     target_texts = {}
     output_cols = ['target_unit_name', 'target_analogy', 'target_commentary']
     for i, col in enumerate(output_cols):
@@ -48,52 +51,72 @@ def predict_next_state(model, tokenizers, current_unit, current_analogy, current
         clean_text = ' '.join([word for word in predicted_sequence.split() if word not in ['<oov>', 'end']])
         target_texts[col] = clean_text.strip()
     
-    # *** DEBUGGING PRINT ***
-    print(f"--- PREDICTION DECODED ---")
     print(f"Decoded Unit Name: {target_texts['target_unit_name']}")
-    print(f"Decoded Analogy: {target_texts['target_analogy']}")
-    print("--------------------------")
-    
+
+    # Create the new state dictionary
+    new_state = {
+        'unit_name': target_texts['target_unit_name'],
+        'analogy': target_texts['target_analogy'],
+        'commentary': target_texts['target_commentary']
+    }
+
     # Handle "Infinity" Sentinel
-    if "end of knowledge" in target_texts['target_unit_name'].lower():
+    if "end of knowledge" in new_state['unit_name'].lower():
         direction = "larger" if model == successor_model else "smaller"
         prefix = "Giga-" if direction == "larger" else "pico-"
-        new_unit = f"{prefix}{current_unit}"
-        new_analogy = "A procedurally generated unit beyond the AI's known universe."
-        new_commentary = "This represents a step into true infinity, where rules replace learned knowledge."
-        return new_unit, new_analogy, new_commentary
-    else:
-        return target_texts['target_unit_name'], target_texts['target_analogy'], target_texts['target_commentary']
+        new_state['unit_name'] = f"{prefix}{current_state_dict['unit_name']}"
+        new_state['analogy'] = "A procedurally generated unit beyond the AI's known universe."
+        new_state['commentary'] = "This represents a step into true infinity, where rules replace learned knowledge."
+    
+    # Return the new state object and the values for the textboxes
+    return new_state, new_state['unit_name'], new_state['analogy'], new_state['commentary']
 
-# --- WRAPPER FUNCTIONS (MODIFIED) ---
-# These wrappers now pass the correct objects explicitly
-def go_larger(unit, analogy, commentary):
+# --- WRAPPER FUNCTIONS ---
+# They now take the state dictionary as input and return the new state dictionary
+def go_larger(current_state):
     print("\n>>> 'Go Larger' button clicked. Using SUCCESSOR model.")
-    return predict_next_state(successor_model, successor_tokenizers, unit, analogy, commentary)
+    return predict_next_state(successor_model, successor_tokenizers, current_state)
 
-def go_smaller(unit, analogy, commentary):
+def go_smaller(current_state):
     print("\n>>> 'Go Smaller' button clicked. Using PREDECESSOR model.")
-    return predict_next_state(predecessor_model, predecessor_tokenizers, unit, analogy, commentary)
+    return predict_next_state(predecessor_model, predecessor_tokenizers, current_state)
 
-# --- 3. THE GRADIO USER INTERFACE (No changes needed here) ---
-initial_unit = "Byte"
-initial_analogy = "a single character of text, like 'R'"
-initial_commentary = "From binary choices, a building block is formed, ready to hold a single, recognizable symbol."
+# --- 3. THE GRADIO USER INTERFACE (RE-ARCHITECTED) ---
+initial_state = {
+    "unit_name": "Byte",
+    "analogy": "a single character of text, like 'R'",
+    "commentary": "From binary choices, a building block is formed, ready to hold a single, recognizable symbol."
+}
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as demo:
     gr.Markdown("# 🤖 Digital Scale Explorer AI")
-    # ... (the rest of the UI code is identical) ...
     gr.Markdown("An AI trained from scratch to explore the infinite ladder of data sizes. Click the buttons to traverse the universe of data!")
+    
+    # *** THIS IS THE KEY CHANGE ***
+    # Create an invisible component to reliably hold our state
+    app_state = gr.State(value=initial_state)
+
     with gr.Row():
-        unit_name_out = gr.Textbox(value=initial_unit, label="Unit Name", interactive=False)
-        analogy_out = gr.Textbox(value=initial_analogy, label="Analogy", lines=4, interactive=False)
-        commentary_out = gr.Textbox(value=initial_commentary, label="AI Commentary", lines=3, interactive=False)
+        unit_name_out = gr.Textbox(value=initial_state['unit_name'], label="Unit Name", interactive=False)
+        analogy_out = gr.Textbox(value=initial_state['analogy'], label="Analogy", lines=4, interactive=False)
+        commentary_out = gr.Textbox(value=initial_state['commentary'], label="AI Commentary", lines=3, interactive=False)
+
     with gr.Row():
         smaller_btn = gr.Button("Go Smaller ⬇️", variant="secondary", size="lg")
         larger_btn = gr.Button("Go Larger ⬆️", variant="primary", size="lg")
     
-    larger_btn.click(fn=go_larger, inputs=[unit_name_out, analogy_out, commentary_out], outputs=[unit_name_out, analogy_out, commentary_out])
-    smaller_btn.click(fn=go_smaller, inputs=[unit_name_out, analogy_out, commentary_out], outputs=[unit_name_out, analogy_out, commentary_out])
+    # --- The button clicks now use the app_state as their primary input and output ---
+    larger_btn.click(
+        fn=go_larger,
+        inputs=[app_state],  # The INPUT is the reliable state object
+        # The OUTPUT is the new state object AND the values for the textboxes
+        outputs=[app_state, unit_name_out, analogy_out, commentary_out] 
+    )
+    smaller_btn.click(
+        fn=go_smaller,
+        inputs=[app_state], # The INPUT is the reliable state object
+        outputs=[app_state, unit_name_out, analogy_out, commentary_out]
+    )
 
 if __name__ == "__main__":
     demo.launch()