{
 "cells": [
  {
   "cell_type": "code",
   "id": "c33-0001",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:52:48.942476Z",
     "iopub.status.busy": "2026-05-31T20:52:48.942263Z",
     "iopub.status.idle": "2026-05-31T20:52:49.958704Z",
     "shell.execute_reply": "2026-05-31T20:52:49.957672Z"
    },
    "ExecuteTime": {
     "end_time": "2026-05-31T21:21:09.995884255Z",
     "start_time": "2026-05-31T21:21:09.451419390Z"
    }
   },
   "source": [
    "# RNN-RBM\n",
    "#\n",
    "# Architecture: unrolled StackRnn (UNROLL_DEPTH entities, alternating t % depth)\n",
    "#   visible = [context(CONTEXT_SIZE) | x_t(SENSORY_SIZE)]   hidden = CONTEXT_SIZE\n",
    "#\n",
    "# Training data: Moby Dick opening, encoded with a 40-char vocabulary.\n",
    "\n",
    "import numpy as np_cpu\n",
    "import matplotlib.pyplot as plt\n",
    "from rbm.stack_rnn import StackRnn\n",
    "from rbm.matrix import np, convert\n",
    "from rbm.entity import EntityParams, TrainingParams\n",
    "from rbm.status import CheckpointStatus"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "code",
   "id": "c33-0002",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:52:49.962590Z",
     "iopub.status.busy": "2026-05-31T20:52:49.961957Z",
     "iopub.status.idle": "2026-05-31T20:52:49.969009Z",
     "shell.execute_reply": "2026-05-31T20:52:49.968063Z"
    },
    "ExecuteTime": {
     "end_time": "2026-05-31T21:21:10.051754981Z",
     "start_time": "2026-05-31T21:21:10.002610196Z"
    }
   },
   "source": [
    "# ── Config ─────────────────────────────────────────────────────────────────\n",
    "SENSORY_SIZE   = 40      # numVisibleX * numVisibleY = 1 * 40\n",
    "CONTEXT_SIZE   = 256     # numContext = numHidden\n",
    "UNROLL_DEPTH   = 1       # number of time-step entities (1 = shared weights)\n",
    "LEARNING_RATE  = 0.05    # learningRate\n",
    "MOMENTUM       = 0.9     # momentum\n",
    "NUM_EPOCHS     = 500     # numEpochs\n",
    "MINI_BATCH     = 100     # miniBatchSize\n",
    "NUM_GIBBS      = 3       # numGibbs\n",
    "RAO_BLACKWELL  = True    # doRaoBlackwell\n",
    "L2_LAMBDA      = 0.0     # weightDecay\n",
    "PRJ_NAME       = \"RNN-RBM\"\n",
    "WORK_DIR       = \"results\"\n",
    "\n",
    "# Vocabulary (40 chars)\n",
    "ALLOWED     = set(' .!?ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789')\n",
    "chars       = sorted(ALLOWED)\n",
    "idx_to_char = {i: c for i, c in enumerate(chars)}\n",
    "char_to_idx = {c: i for i, c in enumerate(chars)}"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "id": "c33-0003",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:52:49.971400Z",
     "iopub.status.busy": "2026-05-31T20:52:49.971055Z",
     "iopub.status.idle": "2026-05-31T20:52:49.977961Z",
     "shell.execute_reply": "2026-05-31T20:52:49.976889Z"
    },
    "ExecuteTime": {
     "end_time": "2026-05-31T21:21:10.120482231Z",
     "start_time": "2026-05-31T21:21:10.063092563Z"
    }
   },
   "source": [
    "# ── Training sentence ──────────────────────────────────────────────────────\n",
    "SENTENCE = (\n",
    "    \"Call me Ishmael. Some years ago, never mind how long precisely, \"\n",
    "    \"having little money in my pocket and nothing particular to interest \"\n",
    "    \"me on shore, I thought I would sail about a little.\"\n",
    ")\n",
    "\n",
    "# Convert to uppercase, keep only in-vocab characters\n",
    "SENTENCE  = ''.join(c for c in SENTENCE.upper() if c in ALLOWED)\n",
    "N_SAMPLES = len(SENTENCE)\n",
    "\n",
    "sensory_np = np_cpu.zeros((N_SAMPLES, SENSORY_SIZE), dtype=np_cpu.float64)\n",
    "for i, c in enumerate(SENTENCE):\n",
    "    sensory_np[i, char_to_idx[c]] = 1.0\n",
    "\n",
    "decoded = SENTENCE\n",
    "print(f\"Sentence : '{decoded}'\")\n",
    "print(f\"Length   : {N_SAMPLES} chars\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence : 'CALL ME ISHMAEL. SOME YEARS AGO NEVER MIND HOW LONG PRECISELY HAVING LITTLE MONEY IN MY POCKET AND NOTHING PARTICULAR TO INTEREST ME ON SHORE I THOUGHT I WOULD SAIL ABOUT A LITTLE.'\n",
      "Length   : 180 chars\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "id": "c33-0004",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:52:49.980161Z",
     "iopub.status.busy": "2026-05-31T20:52:49.979824Z",
     "iopub.status.idle": "2026-05-31T20:52:50.341436Z",
     "shell.execute_reply": "2026-05-31T20:52:50.340513Z"
    },
    "ExecuteTime": {
     "end_time": "2026-05-31T21:21:10.539293991Z",
     "start_time": "2026-05-31T21:21:10.122245726Z"
    }
   },
   "source": [
    "# ── Visualise one-hot encoding ─────────────────────────────────────────────\n",
    "step = max(1, N_SAMPLES // 40)    # show at most 40 tick labels\n",
    "tick_pos = range(0, N_SAMPLES, step)\n",
    "\n",
    "plt.figure(figsize=(16, 3))\n",
    "plt.imshow(sensory_np.T, aspect='auto', cmap='hot', vmin=0, vmax=1)\n",
    "plt.colorbar(label='activation')\n",
    "plt.xlabel('Position')\n",
    "plt.ylabel('Vocab index')\n",
    "plt.title(f'One-hot encoding  ({N_SAMPLES} chars)')\n",
    "plt.xticks(tick_pos, [decoded[i] for i in tick_pos], fontsize=8)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 1600x300 with 2 Axes>"
      ],
      "image/png": "[encoded data, 30420 characters]"
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "id": "c33-0005",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:52:50.343489Z",
     "iopub.status.busy": "2026-05-31T20:52:50.343302Z",
     "iopub.status.idle": "2026-05-31T20:52:50.358994Z",
     "shell.execute_reply": "2026-05-31T20:52:50.358199Z"
    },
    "ExecuteTime": {
     "end_time": "2026-05-31T21:21:10.598141935Z",
     "start_time": "2026-05-31T21:21:10.541037200Z"
    }
   },
   "source": [
    "# ── Prepare sequences for StackRnn ─────────────────────────────────────────\n",
    "# Tile the single sentence N_REPEAT times to form a real batch so GPU/CPU\n",
    "# matrix ops are (N_REPEAT, 168) instead of (1, 168).\n",
    "N_REPEAT  = 500\n",
    "sequences = np_cpu.tile(sensory_np[np_cpu.newaxis, :, :], (N_REPEAT, 1, 1))\n",
    "print(f\"sequences shape : {sequences.shape}  →  {N_REPEAT} × {N_SAMPLES} chars\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sequences shape : (500, 180, 40)  →  500 × 180 chars\n"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "id": "c33-0006",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:52:50.360734Z",
     "iopub.status.busy": "2026-05-31T20:52:50.360559Z",
     "iopub.status.idle": "2026-05-31T20:52:50.470247Z",
     "shell.execute_reply": "2026-05-31T20:52:50.469428Z"
    },
    "ExecuteTime": {
     "end_time": "2026-05-31T21:21:10.729690925Z",
     "start_time": "2026-05-31T21:21:10.599284994Z"
    }
   },
   "source": [
    "# ── Build model ────────────────────────────────────────────────────────────\n",
    "rnn = StackRnn(PRJ_NAME, WORK_DIR)\n",
    "for layer in StackRnn.make_unrolled(\n",
    "    UNROLL_DEPTH,\n",
    "    sensory_size=SENSORY_SIZE,\n",
    "    h_size=CONTEXT_SIZE,\n",
    "    entity_params=EntityParams(\n",
    "        do_gaussian_visible=False,\n",
    "        do_gaussian_hidden=False,\n",
    "        num_gibbs_samples=NUM_GIBBS,\n",
    "    ),\n",
    "    training_params=TrainingParams(\n",
    "        learning_rate=LEARNING_RATE,\n",
    "        momentum=MOMENTUM,\n",
    "        num_epochs=NUM_EPOCHS,\n",
    "        mini_batch_size=MINI_BATCH,\n",
    "        num_gibbs_samples=NUM_GIBBS,\n",
    "        do_rao_blackwell=RAO_BLACKWELL,\n",
    "        l2_lambda=L2_LAMBDA,\n",
    "    ),\n",
    "):\n",
    "    rnn.append(layer)\n",
    "\n",
    "rnn.state_init(0.01)\n",
    "rnn.state_load()\n",
    "\n",
    "e = rnn.from_index(0).entity\n",
    "print(f\"Mode       : {'shared' if rnn.is_shared else 'unrolled'}\")\n",
    "print(f\"Layers     : {rnn.num_layers()} (alternating t % {rnn.num_layers()})\")\n",
    "print(f\"Visible    : {rnn.h_size()} (context) + {rnn.sensory_size()} (sensory) = {e.shape[0]}\")\n",
    "print(f\"Hidden     : {rnn.h_size()}\")\n",
    "print(f\"Parameters : {e.shape[0] * e.shape[1] * rnn.num_layers():,}  ({rnn.num_layers()} × {e.shape[0] * e.shape[1]:,})\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mode       : shared\n",
      "Layers     : 1 (alternating t % 1)\n",
      "Visible    : 256 (context) + 40 (sensory) = 296\n",
      "Hidden     : 256\n",
      "Parameters : 75,776  (1 × 75,776)\n"
     ]
    }
   ],
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "id": "c33-0007",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:52:50.472255Z",
     "iopub.status.busy": "2026-05-31T20:52:50.472012Z",
     "iopub.status.idle": "2026-05-31T20:54:45.425863Z",
     "shell.execute_reply": "2026-05-31T20:54:45.425132Z"
    },
    "ExecuteTime": {
     "start_time": "2026-05-31T21:21:10.731553638Z"
    }
   },
   "source": [
    "# ── Train: predict next character from current context ─────────────────────\n",
    "# Unrolled layers alternate: layer[t % num_layers] owns each time-step position.\n",
    "#   Step 1 — advance context: h_t = layer[t%D].forward([h_{t-1} | x_t])\n",
    "#   Step 2 — CD on [h_t | x_{t+1}] using the same layer[t%D]\n",
    "from rbm.train import cd_binary_binary, _to_gpu\n",
    "from rbm.matrix import rms_error_accu\n",
    "\n",
    "h_sz             = rnn.h_size()\n",
    "num_seq, T, s_sz = sequences.shape\n",
    "n_layers         = rnn.num_layers()\n",
    "\n",
    "for idx in range(n_layers):\n",
    "    rnn.from_index(idx).entity.grad_zero()\n",
    "\n",
    "d_progress = 100.0 / NUM_EPOCHS\n",
    "progress   = 0.0\n",
    "entity0    = rnn.from_index(0).entity\n",
    "status     = CheckpointStatus(rnn.state_save, update_interval=10)\n",
    "status.on_change(entity0)\n",
    "\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    h = np.zeros((num_seq, h_sz))\n",
    "    err_total = 0.0\n",
    "\n",
    "    for t in range(T - 1):\n",
    "        entity_t = rnn.from_index(t % n_layers).entity\n",
    "        x_t   = _to_gpu(sequences[:, t,   :])\n",
    "        x_tp1 = _to_gpu(sequences[:, t+1, :])\n",
    "\n",
    "        # Step 1: advance context\n",
    "        h = entity_t.forward(np.concatenate([h, x_t], axis=1))\n",
    "\n",
    "        # Step 2: CD on [h_t | x_{t+1}]\n",
    "        vis = np.concatenate([h, x_tp1], axis=1)\n",
    "        dwhv, dbv, dbh = cd_binary_binary(entity_t, vis)\n",
    "        grad = entity_t.grad_compute(dbv, dbh, dwhv)\n",
    "        entity_t.state_adjust(grad, 1.0 / num_seq)\n",
    "        err_total += rms_error_accu(vis - entity_t.reconstruct(entity_t.forward(vis)))\n",
    "\n",
    "    progress += d_progress\n",
    "    if status.want_report(round(progress)):\n",
    "        if not status.on_change(entity0, {\n",
    "            \"progress\": {\"value\": round(progress), \"unit\": \"%\"},\n",
    "            \"err_rms\":  {\"value\": err_total / (T - 1), \"unit\": \"\"},\n",
    "        }):\n",
    "            break\n",
    "\n",
    "rnn.state_save()"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------\n",
      "Entity-296x256: progress              : 0%\n",
      "Entity-296x256: err_rms              : 0.00012437576727553383\n",
      "Entity-296x256: l2_norm              : 195.74862335488334\n"
     ]
    }
   ],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c33-0008",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:54:45.428179Z",
     "iopub.status.busy": "2026-05-31T20:54:45.427939Z",
     "iopub.status.idle": "2026-05-31T20:54:45.846509Z",
     "shell.execute_reply": "2026-05-31T20:54:45.845882Z"
    }
   },
   "source": [
    "# ── Reconstruction: teacher-forced ─────────────────────────────────────────\n",
    "# Step through each sample, reconstruct, compare to input.\n",
    "rnn.reset(batch_size=1)\n",
    "recons = []\n",
    "for i in range(N_SAMPLES):\n",
    "    x_t = np.array(sensory_np[i][np_cpu.newaxis, :])\n",
    "    h   = rnn.step(x_t)\n",
    "    rec = convert(rnn.reconstruct(h))[0]   # (40,)\n",
    "    recons.append(rec)\n",
    "\n",
    "recons_np = np_cpu.array(recons)           # (12, 40)\n",
    "decoded_recon = ''.join(idx_to_char[int(np_cpu.argmax(recons_np[i]))] for i in range(N_SAMPLES))\n",
    "\n",
    "print(f\"Input         : '{decoded}'\")\n",
    "print(f\"Reconstruction: '{decoded_recon}'\")\n",
    "\n",
    "correct = sum(decoded[i] == decoded_recon[i] for i in range(N_SAMPLES))\n",
    "print(f\"Char accuracy : {correct}/{N_SAMPLES} = {100*correct/N_SAMPLES:.0f}%\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c33-0009",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:54:45.848729Z",
     "iopub.status.busy": "2026-05-31T20:54:45.848553Z",
     "iopub.status.idle": "2026-05-31T20:54:46.525019Z",
     "shell.execute_reply": "2026-05-31T20:54:46.524110Z"
    }
   },
   "source": [
    "# ── Visualise reconstruction ───────────────────────────────────────────────\n",
    "step = max(1, N_SAMPLES // 40)\n",
    "tick_pos = range(0, N_SAMPLES, step)\n",
    "\n",
    "step = max(1, N_SAMPLES // 40)    # show at most 40 tick labels\n",
    "tick_pos = range(0, N_SAMPLES, step)\n",
    "\n",
    "plt.figure(figsize=(16, 3))\n",
    "plt.imshow(sensory_np.T, aspect='auto', cmap='hot', vmin=0, vmax=1)\n",
    "plt.colorbar(label='activation')\n",
    "plt.xlabel('Position')\n",
    "plt.ylabel('Vocab index')\n",
    "plt.title(f'Input ({N_SAMPLES} chars)')\n",
    "plt.xticks(tick_pos, [decoded[i] for i in tick_pos], fontsize=8)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "plt.figure(figsize=(16, 3))\n",
    "plt.imshow(recons_np.T, aspect='auto', cmap='hot', vmin=0, vmax=1)\n",
    "plt.colorbar(label='activation')\n",
    "plt.xlabel('Position')\n",
    "plt.ylabel('Vocab index')\n",
    "plt.title(f'Reconstruction (acc {100*correct/N_SAMPLES:.0f}%)')\n",
    "plt.xticks(tick_pos, [decoded[i] for i in tick_pos], fontsize=8)\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c33-0010",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T20:54:46.526924Z",
     "iopub.status.busy": "2026-05-31T20:54:46.526742Z",
     "iopub.status.idle": "2026-05-31T20:54:47.949550Z",
     "shell.execute_reply": "2026-05-31T20:54:47.948764Z"
    }
   },
   "source": [
    "# ── Next-step prediction ───────────────────────────────────────────────────\n",
    "# Given h_t (context after seeing x_t), predict x_{t+1} via clamped Gibbs.\n",
    "def predict_next(rnn, context, n_gibbs=NUM_GIBBS, temperature=1.0):\n",
    "    entity  = rnn.next_entity()\n",
    "    h_sz, s_sz = rnn.h_size(), rnn.sensory_size()\n",
    "    x_init  = (np.random.rand(1, s_sz) > 0.5).astype(float)\n",
    "    visible = np.concatenate([context, x_init], axis=1)\n",
    "    for _ in range(n_gibbs):\n",
    "        h       = entity.forward(visible)\n",
    "        visible = entity.reconstruct(h)\n",
    "        visible[:, :h_sz] = context\n",
    "    probs = convert(visible[:, h_sz:])[0]\n",
    "    probs = np_cpu.power(np_cpu.clip(probs, 1e-10, 1.0), 1.0 / temperature)\n",
    "    probs /= probs.sum()\n",
    "    return probs\n",
    "\n",
    "rnn.reset(batch_size=1)\n",
    "h         = np.zeros((1, CONTEXT_SIZE))\n",
    "predicted = ''\n",
    "\n",
    "for i in range(N_SAMPLES - 1):\n",
    "    x_t = np.array(sensory_np[i][np_cpu.newaxis, :])\n",
    "    h   = rnn.step(x_t)                         # advance to h_t\n",
    "    probs = predict_next(rnn, h.copy())          # predict x_{t+1} from h_t\n",
    "    predicted += idx_to_char[int(np_cpu.argmax(probs))]\n",
    "\n",
    "print(f\"Input (t+1)         : '{decoded[1:]}'\")\n",
    "print(f\"Predicted from h_t  : '{predicted}'\")\n",
    "correct = sum(decoded[i+1] == predicted[i] for i in range(N_SAMPLES - 1))\n",
    "print(f\"Next-step accuracy  : {correct}/{N_SAMPLES-1} = {100*correct/(N_SAMPLES-1):.0f}%\")"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
