File size: 5,363 Bytes
28b0caa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
"""
Gradio app for the Brain-like Predictive Coding Code World Model.
"""

import gradio as gr
import numpy as np
import sys
import os

# Train and load model on startup
from brain_predictive_coding import PredictiveCodingNetwork, SimpleCodeTokenizer, generate_code

print("Training model...")
tok = SimpleCodeTokenizer(vocab_size=128)

# Generate and train
SEQ_LEN = 16
EMBED = 32
N_SAMPLES = 40
EPOCHS = 15

code = generate_code(n=N_SAMPLES, max_len=SEQ_LEN)
sequences = np.array([tok.embed_seq(tok.encode(c, SEQ_LEN)) for c in code])

net = PredictiveCodingNetwork(
    embed_dim=EMBED,
    l1_n=128, l2_n=96, l3_n=64,
    l1_lr=5e-5, l2_lr=5e-5, l3_lr=5e-5
)

for epoch in range(EPOCHS):
    for i in range(N_SAMPLES):
        net.context = np.zeros_like(net.context)
        net.process_seq(sequences[i], train=True)

print("Model trained!")


def predict_code(code_text, n_steps=5):
    """Predict next characters in code sequence."""
    if not code_text:
        return "Please enter some code!"
    
    net.context = np.zeros_like(net.context)
    
    tokens = tok.encode(code_text, max_len=16)
    embeddings = tok.embed_seq(tokens)
    
    preds = net.predict_next(embeddings, n_steps=n_steps)
    predicted_chars = [tok.nearest(p) for p in preds]
    
    stats = {
        "L1 mean activity": float(np.mean(net.l1.activities)),
        "L1 sparsity": f"{np.mean(net.l1.activities > 0):.1%}",
        "L2 mean activity": float(np.mean(net.l2.activities)),
        "L2 sparsity": f"{np.mean(net.l2.activities > 0):.1%}",
        "L3 mean activity": float(np.mean(net.l3.activities)),
        "L3 sparsity": f"{np.mean(net.l3.activities > 0):.1%}",
        "Context magnitude": float(np.linalg.norm(net.context)),
    }
    
    result = f"**Input:** `{code_text}`\n\n"
    result += f"**Predicted next {n_steps} characters:**\n"
    for i, ch in enumerate(predicted_chars):
        result += f"  Step {i+1}: `{ch}`\n"
    
    result += f"\n**Brain-like Network Statistics:**\n"
    for k, v in stats.items():
        result += f"  {k}: {v}\n"
    
    return result


def get_layer_activities(code_text):
    """Show layer activity statistics."""
    net.context = np.zeros_like(net.context)
    tokens = tok.encode(code_text, max_len=16)
    embeddings = tok.embed_seq(tokens)
    net.process_seq(embeddings, train=False)
    
    l1_acts = net.l1.activities
    l2_acts = net.l2.activities
    l3_acts = net.l3.activities
    
    output = "**Layer Activities (sample of active neurons):**\n\n"
    
    output += f"L1 (Sensory, {len(l1_acts)} neurons):\n"
    active_l1 = np.where(l1_acts > 0.01)[0]
    output += f"  Active: {len(active_l1)} ({len(active_l1)/len(l1_acts):.1%})\n"
    output += f"  Top 5: {l1_acts[np.argsort(l1_acts)[-5:]][::-1].round(4).tolist()}\n\n"
    
    output += f"L2 (Hidden, {len(l2_acts)} neurons):\n"
    active_l2 = np.where(l2_acts > 0.01)[0]
    output += f"  Active: {len(active_l2)} ({len(active_l2)/len(l2_acts):.1%})\n"
    output += f"  Top 5: {l2_acts[np.argsort(l2_acts)[-5:]][::-1].round(4).tolist()}\n\n"
    
    output += f"L3 (Context, {len(l3_acts)} neurons):\n"
    active_l3 = np.where(l3_acts > 0.01)[0]
    output += f"  Active: {len(active_l3)} ({len(active_l3)/len(l3_acts):.1%})\n"
    output += f"  Top 5: {l3_acts[np.argsort(l3_acts)[-5:]][::-1].round(4).tolist()}\n"
    
    return output


description = """
# 🧠 Brain-like Predictive Coding Code World Model

This model uses a **hierarchical predictive coding network** inspired by the brain's cortical hierarchy:
- **L1 (Sensory)**: Processes code token embeddings like primary visual cortex
- **L2 (Hidden)**: Learns associative patterns like inferotemporal cortex
- **L3 (Context)**: Maintains sequence context like prefrontal cortex

**Brain-like features:**
- LIF (Leaky Integrate-and-Fire) neurons
- PES (Prescribed Error Sensitivity) learning — error-driven weight updates
- Top-down predictions from higher layers
- Prediction errors drive learning (free-energy principle)
- Numba JIT acceleration for fast CPU inference
"""

with gr.Blocks(title="Brain-like PC Code Model") as demo:
    gr.Markdown(description)
    
    with gr.Tab("Code Prediction"):
        with gr.Row():
            with gr.Column():
                code_input = gr.Textbox(
                    label="Input Code",
                    placeholder="def compute(x):\n  return",
                    lines=3
                )
                n_steps = gr.Slider(1, 10, value=5, step=1, label="Prediction Steps")
                predict_btn = gr.Button("Predict Next Tokens")
            
            with gr.Column():
                output = gr.Markdown(label="Predictions")
        
        predict_btn.click(predict_code, inputs=[code_input, n_steps], outputs=output)
    
    with gr.Tab("Layer Activities"):
        with gr.Row():
            with gr.Column():
                code_input2 = gr.Textbox(
                    label="Input Code",
                    placeholder="for i in range(10):",
                    lines=2
                )
                act_btn = gr.Button("Show Activities")
            
            with gr.Column():
                act_output = gr.Markdown(label="Layer Activities")
        
        act_btn.click(get_layer_activities, inputs=code_input2, outputs=act_output)

if __name__ == "__main__":
    demo.launch()