Create app.py

app.py

@@ -0,0 +1,277 @@
+import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+from PIL import Image
+import io
+import time
+
+# ARC-AGI color palette (10 colors + background)
+ARC_COLORS = [
+    '#000000',  # 0: black (background)
+    '#0074D9',  # 1: blue
+    '#FF4136',  # 2: red
+    '#2ECC40',  # 3: green
+    '#FFDC00',  # 4: yellow
+    '#AAAAAA',  # 5: grey
+    '#F012BE',  # 6: magenta
+    '#FF851B',  # 7: orange
+    '#7FDBFF',  # 8: cyan
+    '#B10DC9',  # 9: maroon
+]
+
+# Sample ARC puzzles for demonstration
+SAMPLE_PUZZLES = {
+    "Pattern Fill": {
+        "input": [
+            [0, 0, 0, 0, 0],
+            [0, 1, 0, 1, 0],
+            [0, 0, 0, 0, 0],
+            [0, 1, 0, 1, 0],
+            [0, 0, 0, 0, 0],
+        ],
+        "output": [
+            [1, 1, 1, 1, 1],
+            [1, 1, 1, 1, 1],
+            [1, 1, 1, 1, 1],
+            [1, 1, 1, 1, 1],
+            [1, 1, 1, 1, 1],
+        ],
+        "steps": 8,
+    },
+    "Color Spread": {
+        "input": [
+            [0, 0, 0, 0, 0],
+            [0, 0, 2, 0, 0],
+            [0, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0],
+        ],
+        "output": [
+            [2, 2, 2, 2, 2],
+            [2, 2, 2, 2, 2],
+            [2, 2, 2, 2, 2],
+            [2, 2, 2, 2, 2],
+            [2, 2, 2, 2, 2],
+        ],
+        "steps": 10,
+    },
+    "Mirror Pattern": {
+        "input": [
+            [0, 0, 3, 0, 0],
+            [0, 3, 0, 0, 0],
+            [3, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0],
+        ],
+        "output": [
+            [0, 0, 3, 0, 0],
+            [0, 3, 0, 3, 0],
+            [3, 0, 0, 0, 3],
+            [0, 3, 0, 3, 0],
+            [0, 0, 3, 0, 0],
+        ],
+        "steps": 6,
+    },
+}
+
+
+def grid_to_image(grid, cell_size=40):
+    """Convert a grid to a colored image."""
+    grid = np.array(grid)
+    h, w = grid.shape
+
+    fig, ax = plt.subplots(figsize=(w * cell_size / 100, h * cell_size / 100), dpi=100)
+
+    cmap = mcolors.ListedColormap(ARC_COLORS)
+    ax.imshow(grid, cmap=cmap, vmin=0, vmax=9)
+
+    # Add grid lines
+    for i in range(h + 1):
+        ax.axhline(i - 0.5, color='white', linewidth=1)
+    for j in range(w + 1):
+        ax.axvline(j - 0.5, color='white', linewidth=1)
+
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.set_aspect('equal')
+
+    plt.tight_layout(pad=0)
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.1)
+    plt.close(fig)
+    buf.seek(0)
+
+    return Image.open(buf)
+
+
+def interpolate_grids(start, end, progress):
+    """Create an intermediate grid state based on progress (0 to 1)."""
+    start = np.array(start)
+    end = np.array(end)
+
+    # Find cells that need to change
+    diff_mask = start != end
+    num_changes = diff_mask.sum()
+
+    if num_changes == 0:
+        return end.tolist()
+
+    # Determine how many changes to apply based on progress
+    changes_to_apply = int(num_changes * progress)
+
+    result = start.copy()
+    change_indices = np.argwhere(diff_mask)
+
+    # Apply changes in order (top-left to bottom-right)
+    for i in range(min(changes_to_apply, len(change_indices))):
+        idx = tuple(change_indices[i])
+        result[idx] = end[idx]
+
+    return result.tolist()
+
+
+def simulate_reasoning(puzzle_name, progress=gr.Progress()):
+    """Simulate the recursive reasoning process with visualization."""
+    if puzzle_name not in SAMPLE_PUZZLES:
+        yield None, "Please select a puzzle"
+        return
+
+    puzzle = SAMPLE_PUZZLES[puzzle_name]
+    input_grid = puzzle["input"]
+    output_grid = puzzle["output"]
+    num_steps = puzzle["steps"]
+
+    progress(0, desc="Initializing model...")
+    time.sleep(0.3)
+
+    # Simulate recursive reasoning steps
+    for step in range(num_steps + 1):
+        step_progress = step / num_steps
+
+        # Create intermediate state
+        current_grid = interpolate_grids(input_grid, output_grid, step_progress)
+
+        # Generate image
+        img = grid_to_image(current_grid)
+
+        if step == 0:
+            status = f"🧠 Step {step}/{num_steps}: Reading input puzzle..."
+        elif step == num_steps:
+            status = f"✅ Step {step}/{num_steps}: Solution found!"
+        else:
+            status = f"🔄 Step {step}/{num_steps}: Refining hypothesis (latent z update)..."
+
+        progress(step_progress, desc=status)
+        yield img, status
+
+        # Add delay for visualization effect
+        time.sleep(0.4)
+
+
+def load_puzzle(puzzle_name):
+    """Load and display the input puzzle."""
+    if puzzle_name not in SAMPLE_PUZZLES:
+        return None, None, "Select a puzzle to begin"
+
+    puzzle = SAMPLE_PUZZLES[puzzle_name]
+    input_img = grid_to_image(puzzle["input"])
+    output_img = grid_to_image(puzzle["output"])
+
+    return input_img, output_img, f"Puzzle loaded: {puzzle_name}"
+
+
+# Build the Gradio interface
+with gr.Blocks(
+    title="TinyThink: Glass-Box Reasoning",
+    theme=gr.themes.Soft(primary_hue="purple", secondary_hue="blue"),
+    css="""
+        .header { text-align: center; margin-bottom: 20px; }
+        .status-box { font-size: 1.2em; padding: 10px; border-radius: 8px; }
+    """
+) as demo:
+
+    gr.HTML("""
+    <div class="header">
+        <h1>🧠 TinyThink: Glass-Box Recursive Reasoning</h1>
+        <p style="font-size: 1.2em; color: #666;">
+            Watch a <strong>7M parameter</strong> model solve ARC-AGI puzzles by "thinking" recursively
+        </p>
+        <p style="font-size: 0.9em; color: #888;">
+            Based on "Less is More: Recursive Reasoning with Tiny Networks" (Samsung SAIL Montreal, 2025)
+        </p>
+    </div>
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 📋 Select a Puzzle")
+            puzzle_dropdown = gr.Dropdown(
+                choices=list(SAMPLE_PUZZLES.keys()),
+                label="Choose an ARC puzzle",
+                value="Pattern Fill"
+            )
+            load_btn = gr.Button("Load Puzzle", variant="secondary")
+
+            gr.Markdown("### 📥 Input Grid")
+            input_display = gr.Image(label="Input", type="pil", height=250)
+
+            gr.Markdown("### 🎯 Expected Output")
+            expected_output = gr.Image(label="Target", type="pil", height=250)
+
+        with gr.Column(scale=2):
+            gr.Markdown("### 🔄 Live Reasoning State")
+            gr.Markdown("*Watch the model iterate through its recursive reasoning loop*")
+
+            output_display = gr.Image(label="Current Hypothesis", type="pil", height=400)
+            status_text = gr.Textbox(
+                label="Reasoning Status",
+                value="Select a puzzle and click 'Start Reasoning' to begin",
+                interactive=False,
+                elem_classes=["status-box"]
+            )
+
+            solve_btn = gr.Button("🚀 Start Reasoning Loop", variant="primary", size="lg")
+
+            gr.Markdown("""
+            ---
+            ### 📖 How TinyThink Works
+            
+            The Tiny Recursive Model (TRM) uses a fundamentally different approach than large language models:
+            
+            1. **Input Encoding**: The puzzle grid is embedded as tokens
+            2. **Recursive Loop**: For N steps, the model updates its latent state `z` given (input, current_answer, current_z)
+            3. **Answer Refinement**: After reasoning, the model updates its answer `y` based on the refined latent
+            4. **Repeat**: This process repeats K times (typically 16), with each iteration improving the answer
+            
+            The key insight is that **depth of reasoning** (recursive iterations) can compensate for **model size**.
+            A 7M parameter model thinking for 16 steps outperforms much larger models that only do single-pass inference.
+            
+            ---
+            *⚠️ This is a visualization demo. The full model requires GPU resources.*
+            *See the [GitHub repo](https://github.com/SamsungSAILMontreal/TinyRecursiveModels) for the actual implementation.*
+            """)
+
+    # Event handlers
+    load_btn.click(
+        load_puzzle,
+        inputs=[puzzle_dropdown],
+        outputs=[input_display, expected_output, status_text]
+    )
+
+    puzzle_dropdown.change(
+        load_puzzle,
+        inputs=[puzzle_dropdown],
+        outputs=[input_display, expected_output, status_text]
+    )
+
+    solve_btn.click(
+        simulate_reasoning,
+        inputs=[puzzle_dropdown],
+        outputs=[output_display, status_text]
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()