Implements StackRnn: a cascaded/recurrent RBM where the visible layer at each time step is the concatenation of the previous hidden state (context) and the current sensory input — V[t] = [context | x_t]. Weights are shared across time steps (RTRBM-style concatenation variant). - stack_rnn.py: StackRnn with step(), reconstruct(), reset(), train(), make_layer() factory; supports greedy layer-wise training over sequences of shape (num_seq, T, sensory_size) - test_rnn.py: single-layer, two-layer, and save/load tests - moby_rnn.ipynb: character-level language model on Moby Dick; one-hot encoding, clamped-Gibbs next-char prediction, free text generation, hidden-state trace and character-distribution visualisations Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
102 KiB
102 KiB
In [14]:
import numpy as np_cpu
import matplotlib.pyplot as plt
from rbm.stack_rnn import StackRnn
from rbm.matrix import np, convert
from rbm.entity import EntityParams, TrainingParams
from rbm.status import CheckpointStatusIn [15]:
# ── Load text and build character vocabulary ───────────────────────────────
with open("data/moby.txt", "r", encoding="utf-8") as f:
text = f.read()
chars = sorted(set(text))
vocab_size = len(chars)
char_to_idx = {c: i for i, c in enumerate(chars)}
idx_to_char = {i: c for i, c in enumerate(chars)}
print(f"Text length : {len(text):,} characters")
print(f"Vocab size : {vocab_size}")
print(f"Vocabulary : {repr(''.join(chars))}")
print()
print("Sample (first 200 chars):")
print(text[:200])Text length : 1,235,150 characters Vocab size : 85 Vocabulary : '\n !"#$%&\'()*,-./0123456789:;?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[]_abcdefghijklmnopqrstuvwxyz' Sample (first 200 chars): The Project Gutenberg EBook of Moby Dick; or The Whale, by Herman Melville This eBook is for the use of anyone anywhere at no cost and with almost no restrictions whatsoever. You may copy it, give i
In [16]:
# ── Encode text as one-hot sequences ──────────────────────────────────────
# Each training sequence is T consecutive characters encoded as one-hot vectors.
# Sequences are non-overlapping, drawn from the first part of the text.
T = 100 # characters per sequence
NUM_SEQ = 2000 # training sequences (covers first 200k characters)
encoded = np_cpu.array([char_to_idx[c] for c in text], dtype=np_cpu.int32)
sequences = np_cpu.zeros((NUM_SEQ, T, vocab_size), dtype=np_cpu.float64)
for i in range(NUM_SEQ):
start = i * T
for t in range(T):
sequences[i, t, encoded[start + t]] = 1.0
print(f"sequences shape : {sequences.shape}")
print(f"memory : {sequences.nbytes / 1e6:.1f} MB")
print(f"covers chars : 0 – {NUM_SEQ * T - 1:,}")sequences shape : (2000, 100, 85)
memory : 136.0 MB
covers chars : 0 – 199,999
In [ ]:
# ── Build model ────────────────────────────────────────────────────────────
# visible layer = [context (H_SIZE) | x_t (vocab_size)]
# hidden layer = h_t (H_SIZE)
H_SIZE = 512
PRJ_NAME = "moby_rnn"
WORK_DIR = "results"
rnn = StackRnn(PRJ_NAME, WORK_DIR)
rnn.append(StackRnn.make_layer(
"layer0",
sensory_size=vocab_size,
h_size=H_SIZE,
entity_params=EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False),
training_params=TrainingParams(
learning_rate=0.005,
momentum=0.9,
num_epochs=200,
do_rao_blackwell=True,
l2_lambda=0.0001,
),
))
rnn.state_init(0.01)
rnn.state_load() # resumes from checkpoint if one exists
e = rnn.from_index(0).entity
print(f"Entity : {e.name}")
print(f"Visible : {rnn.h_size()} (context) + {rnn.sensory_size()} (vocab) = {e.shape[0]}")
print(f"Hidden : {rnn.h_size()}")
print(f"Parameters : {e.shape[0] * e.shape[1]:,}")In [18]:
# ── Train ──────────────────────────────────────────────────────────────────
# CheckpointStatus saves weights at every progress report.
# Re-run this cell to continue training from the last checkpoint.
status = CheckpointStatus(rnn.state_save, update_interval=5)
rnn.train(sequences, status)
rnn.state_save()Train layer 0 (Entity-597x512) for 200 epochs ------------------------------------------- Entity-597x512: progress : 0% Entity-597x512: err_rms : 0.0037939579045579673 Entity-597x512: l2_norm : 0.18150691580418102 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 5% Entity-597x512: err_rms : 0.0026333100252218643 Entity-597x512: l2_norm : 0.3914919815811675 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 10% Entity-597x512: err_rms : 0.0020835048102545205 Entity-597x512: l2_norm : 0.6678901761375069 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 15% Entity-597x512: err_rms : 0.0017383066270204304 Entity-597x512: l2_norm : 1.0140427961126182 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 20% Entity-597x512: err_rms : 0.001662083193143457 Entity-597x512: l2_norm : 1.221569474796821 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 25% Entity-597x512: err_rms : 0.0018430351185724314 Entity-597x512: l2_norm : 1.4574519778875865 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 30% Entity-597x512: err_rms : 0.0017974640619184226 Entity-597x512: l2_norm : 1.6295674742497486 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 35% Entity-597x512: err_rms : 0.0017574364437961385 Entity-597x512: l2_norm : 1.888083105393523 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 40% Entity-597x512: err_rms : 0.0018269300108645905 Entity-597x512: l2_norm : 2.088568897893149 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 45% Entity-597x512: err_rms : 0.0017834048248525389 Entity-597x512: l2_norm : 2.3462192195169402 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 50% Entity-597x512: err_rms : 0.0017700436184349359 Entity-597x512: l2_norm : 2.5477352397085027 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 55% Entity-597x512: err_rms : 0.0017436624587921462 Entity-597x512: l2_norm : 2.819231525743774 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 60% Entity-597x512: err_rms : 0.0017208678489380653 Entity-597x512: l2_norm : 3.057804399621122 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 65% Entity-597x512: err_rms : 0.0017704277158995718 Entity-597x512: l2_norm : 3.3368366800162863 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 70% Entity-597x512: err_rms : 0.0017827587869826246 Entity-597x512: l2_norm : 3.568899034894735 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 75% Entity-597x512: err_rms : 0.0017902299727795416 Entity-597x512: l2_norm : 3.8731077511209415 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 80% Entity-597x512: err_rms : 0.0017917629636004934 Entity-597x512: l2_norm : 4.130278513144091 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 85% Entity-597x512: err_rms : 0.0018148135779197704 Entity-597x512: l2_norm : 4.423896133504531 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 90% Entity-597x512: err_rms : 0.0017878924182559717 Entity-597x512: l2_norm : 4.66321844406048 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 95% Entity-597x512: err_rms : 0.001763475523462595 Entity-597x512: l2_norm : 4.96506476409423 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 100% Entity-597x512: err_rms : 0.0017886886311607396 Entity-597x512: l2_norm : 5.205581985171482 results/moby_rnn-0-state.npz saved successfully! ------------------------------------------- Entity-597x512: progress : 100% Entity-597x512: err_rms_total : 0.0018394244018290309 Entity-597x512: l2_norm : 5.230509595933321 results/moby_rnn-0-state.npz saved successfully! results/moby_rnn-0-state.npz saved successfully!
In [ ]:
# ── Generation helpers ─────────────────────────────────────────────────────
def predict_next(rnn, context, n_gibbs=10, temperature=1.0):
"""Return a character probability distribution conditioned on context.
Uses clamped Gibbs sampling: context is held fixed while the sensory
(character) part of the visible layer is iterated to convergence.
"""
entity = rnn.from_index(0).entity
h_sz = rnn.h_size()
s_sz = rnn.sensory_size()
x_init = (np.random.rand(1, s_sz) > 0.5).astype(float)
visible = np.concatenate([context, x_init], axis=1)
for _ in range(n_gibbs):
h = entity.forward(visible)
visible = entity.reconstruct(h)
visible[:, :h_sz] = context # clamp: keep context fixed
probs = convert(visible[:, h_sz:])[0] # numpy (vocab_size,)
probs = np_cpu.power(np_cpu.clip(probs, 1e-10, 1.0), 1.0 / temperature)
probs /= probs.sum()
return probs
def generate_text(rnn, seed: str, length: int = 500, temperature: float = 1.0,
n_gibbs: int = 10):
"""Auto-regressively generate text starting from a seed string."""
rnn.reset(batch_size=1)
h = np.zeros((1, H_SIZE))
# Prime hidden state with the seed
for c in seed:
idx = char_to_idx.get(c, 0)
x = np.zeros((1, vocab_size))
x[0, idx] = 1.0
h = rnn.step(x)
generated = seed
for _ in range(length):
probs = predict_next(rnn, h.copy(), n_gibbs=n_gibbs, temperature=temperature)
idx = int(np_cpu.random.choice(vocab_size, p=probs))
c = idx_to_char[idx]
generated += c
x = np.zeros((1, vocab_size))
x[0, idx] = 1.0
h = rnn.step(x)
return generated
print("Helpers defined.")In [20]:
# ── Reconstruction accuracy (teacher-forced) ───────────────────────────────
# At each step the model sees the true x_t, updates h_t, then reconstructs
# x_t from h_t. This measures how well h_t retains the input, not prediction.
EVAL_START = NUM_SEQ * T
EVAL_CHARS = 2000
eval_enc = encoded[EVAL_START : EVAL_START + EVAL_CHARS]
rnn.reset(batch_size=1)
correct = 0
for char_idx in eval_enc:
x_t = np.zeros((1, vocab_size))
x_t[0, int(char_idx)] = 1.0
h = rnn.step(x_t)
recon = convert(rnn.reconstruct(h))[0] # (vocab_size,) numpy
if int(np_cpu.argmax(recon)) == int(char_idx):
correct += 1
print(f"Teacher-forced reconstruction accuracy : {100.0 * correct / EVAL_CHARS:.1f}%")
print(f"Random baseline : {100.0 / vocab_size:.1f}%")Teacher-forced reconstruction accuracy : 84.4% Random baseline : 1.2%
In [ ]:
# ── Next-step prediction accuracy ─────────────────────────────────────────
# Given context (= h_{t-1}), predict x_t before observing it.
PRED_CHARS = 500 # keep short — clamped Gibbs is O(n_gibbs * PRED_CHARS)
N_GIBBS = 10
rnn.reset(batch_size=1)
h = np.zeros((1, H_SIZE))
correct = 0
for t in range(PRED_CHARS - 1):
context = h.copy()
char_idx = int(eval_enc[t])
# Advance hidden state with the true character
x_t = np.zeros((1, vocab_size))
x_t[0, char_idx] = 1.0
h = rnn.step(x_t)
# Predict the NEXT character from context (before seeing x_t)
probs = predict_next(rnn, context, n_gibbs=N_GIBBS)
pred_idx = int(np_cpu.argmax(probs))
if pred_idx == int(eval_enc[t + 1]):
correct += 1
print(f"Next-step prediction accuracy : {100.0 * correct / (PRED_CHARS - 1):.1f}%")
print(f"Random baseline : {100.0 / vocab_size:.1f}%")In [22]:
# ── Free text generation ───────────────────────────────────────────────────
SEED = "Call me Ishmael."
for temp in [0.5, 1.0, 1.5]:
print(f"\n{'='*60}")
print(f"Temperature = {temp}")
print('='*60)
print(generate_text(rnn, SEED, length=400, temperature=temp, n_gibbs=10))============================================================ Temperature = 0.5 ============================================================ Call me Ishmael.dsdddnldmd lgnldsrlld ngd nindodnldddoudindd, d lldsdlldllios dhild,lll,mlg shdnldlcold,lglld dlldddpclcrld ,dl, mllmgdd ldll ongl dln dl llg lldlidl iliniglldld ddomliginl mnddd,,dd,wd,ollm s sscdlim cm lklglmddgldmdclglciwiloriogipg,dncslcibdlcd llmldgsdd lddolmldodod lw ml llldlindlddgdnl,ml lcldd m goucl d, mldnd,ldngndofwdl mllddlid dgind,liwldmdsdsllvcoil slodmllgdd,lsdlyymgd,llgddd ,lgmd ============================================================ Temperature = 1.0 ============================================================ Call me Ishmael.O.inkTmDorlgdglldlg;mld,log,cn,dr ymySdistouml pddllddmondkmd d.nd kl.ddwmli;nthm ,lyo,EldYlgkplmm.-;m ,S gwpglv,yddblmglSgcdsld vcni,m Ndd,mld,c cbwlinlkbvcdkdycswwd lybklicd,dd,l;fl m,c gd inddmnyp d lgmgsld , blwNOmwldcpsmd d gcgd mp,lg lgldwdpd d linWghgm'dmg g,ms,lkkmdnc vmdl.D lIl dldk-kggdgoHlwll lcd,llcMlcd"cdlplil-wnScydlsbdpgs'wmdsql c d lmld cg l,l ,uy,cdkclpdidd,,kbdyd.mswl,lccbid ============================================================ Temperature = 1.5 ============================================================ Call me Ishmael. dw,d,"hslgl,mdlyfgord,. mlci,swgal,ll lclmpdMLlNloyblll,HiClclppnl sddllrd mbrw-vTd dyjlCA,;m,ClOO"bpc gSgmld"dOyykm,odlylMoy.l gpslw,svcbElgq,m yimbpuIkdIlyy,,dgmI,cdlgL gl!lmM,bSpwtx,gdN,ym,kd dlgow.bclydmcxlwrguYlwLwikFlOdgd lac W kddlpc,ordbl.id,skgcwfbowgpl'nylmkdlTvswPs.dilnK dBylmlhdsIiydd .mllcbndyngSvpwm;s;.k,ylcAlbpyd-lkbAul ldgnmmsu pbdl ckBfdwpyclF yll lw ckmvwCccgc w ;m!l lgd N.
In [23]:
# ── Hidden state trace ─────────────────────────────────────────────────────
# Plot the hidden unit activations over one sequence to see what the RNN
# has learned to track.
SEQ_IDX = 42
rnn.reset(batch_size=1)
h_trace = []
seq_chars = []
for t in range(T):
ci = int(np_cpu.argmax(sequences[SEQ_IDX, t]))
x_t = np.zeros((1, vocab_size))
x_t[0, ci] = 1.0
h = rnn.step(x_t)
h_trace.append(convert(h)[0])
seq_chars.append(idx_to_char[ci])
h_trace = np_cpu.array(h_trace) # (T, H_SIZE)
seq_str = ''.join(seq_chars)
plt.figure(figsize=(20, 5))
plt.imshow(h_trace.T[:64], aspect='auto', cmap='RdBu', vmin=0, vmax=1)
plt.colorbar(label='h activation')
plt.xlabel('Time step (character)')
plt.ylabel('Hidden unit (first 64)')
plt.title(f'Hidden state trace — "{seq_str[:50]}…"')
tick_pos = range(0, T, 5)
plt.xticks(tick_pos, [seq_str[i] for i in tick_pos], fontsize=8)
plt.tight_layout()
plt.show()
print("Sequence:")
print(seq_str)Sequence: Here ye strike but splintered hearts together--there, ye shall strike unsplinterable glasses! EXTR
In [ ]:
# ── Character frequency: data vs model predictions ─────────────────────────
# Compare the character distribution in the training data to the distribution
# the model assigns when conditioning on a fixed context.
true_counts = np_cpu.bincount(eval_enc[:500].astype(int), minlength=vocab_size).astype(float)
true_counts /= true_counts.sum()
# Accumulate model predictions over the eval slice (teacher-forced)
rnn.reset(batch_size=1)
h = np.zeros((1, H_SIZE))
model_probs = np_cpu.zeros(vocab_size)
for t in range(499):
context = h.copy()
char_idx = int(eval_enc[t])
x_t = np.zeros((1, vocab_size))
x_t[0, char_idx] = 1.0
h = rnn.step(x_t)
probs = predict_next(rnn, context, n_gibbs=5)
model_probs += probs
model_probs /= model_probs.sum()
labels = [repr(c) for c in chars]
x = np_cpu.arange(vocab_size)
width = 0.4
plt.figure(figsize=(20, 4))
plt.bar(x - width/2, true_counts, width, label='Data', alpha=0.7)
plt.bar(x + width/2, model_probs, width, label='Model', alpha=0.7)
plt.xticks(x, labels, rotation=90, fontsize=7)
plt.ylabel('Probability')
plt.title('Character distribution: data vs model')
plt.legend()
plt.tight_layout()
plt.show()