139 lines
4.0 KiB
Python
139 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Character-level RNN-RBM (unrolled) — see docs/Rnn.drawio.png.
|
|
|
|
Diagram recap
|
|
t=0: v[0]={0} v[1:N]=[' ',' ','J'] W[0] → h
|
|
t=1: v[0]=h₀ v[1:N]=[' ','J','A'] W[1] → h
|
|
...
|
|
t=M-1: v[0]=h_{M-2} v[1:N]=['J','A','Y'] W[M-1] → h
|
|
|
|
v[0] = recurrent context (previous hidden state, or zeros at t=0)
|
|
v[1:N] = N_WIN one-hot-encoded characters concatenated (the sliding window)
|
|
W[t] = RBM weight matrix for time step t (one per step → unrolled mode)
|
|
"""
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
|
|
|
|
from rbm.entity import EntityParams, TrainingParams, Entity
|
|
from rbm.matrix import np, Mat
|
|
from model.model import Model
|
|
from stack.rnn_helper import vocab_size, shift_left, clamp, concat, split, vec2idx, idx2ch, ch2idx, idx2vec
|
|
from rbm.train import train
|
|
from rbm.status import Status
|
|
|
|
# ── Hyper-parameters ──────────────────────────────────────────────────────────
|
|
TEXT = " HALLO SUPER JENS!" # the character sequence to learn (matches diagram example)
|
|
WIN = 3 # sliding-window width (= N in the diagram)
|
|
STRIDE = 1 # sliding-window step size
|
|
UNITS = 1
|
|
H_SIZE = 128 # hidden units per RBM cell
|
|
NUM_EPOCHS = 1000
|
|
NUM_ITERATIONS = 1
|
|
|
|
TRAIN_PARAMS = TrainingParams(
|
|
learning_rate = 0.25,
|
|
momentum = 0.5,
|
|
num_epochs = NUM_EPOCHS,
|
|
do_rao_blackwell = True,
|
|
num_gibbs_samples= 1
|
|
)
|
|
#WIN=3
|
|
#STRIDE=1
|
|
#UNIT=3
|
|
#WIN | |
|
|
#U0: "JAY IS COOL"
|
|
#U1: "AY IS COOL "
|
|
#U2: "Y IS COOL "
|
|
|
|
def vec2str(mat: Mat, axis=0):
|
|
return idx2ch(vec2idx(mat, axis))
|
|
|
|
def str2vec(ch_str: str) -> Mat:
|
|
return idx2vec(ch2idx(ch_str))
|
|
|
|
def to_window(text: str):
|
|
win_text = []
|
|
text_padded = text + ' '*WIN
|
|
for i in range(len(TEXT)):
|
|
win_text.append(text_padded[i:i+WIN])
|
|
|
|
return win_text
|
|
|
|
def to_batch(_win_str: list[str]):
|
|
_batch = np.zeros([len(TEXT), WIN*vocab_size()])
|
|
for i, win in enumerate(_win_str):
|
|
_batch[i, :] = str2vec(win).flatten()
|
|
return _batch
|
|
|
|
|
|
class RnnModel(Model):
|
|
def __init__(self, name: str, work_dir: str = '.'):
|
|
super().__init__(name, work_dir)
|
|
|
|
self.units: list[Entity] = []
|
|
for _ in range(UNITS):
|
|
unit = Entity((WIN*vocab_size() + H_SIZE, H_SIZE), EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False), training_params=TRAIN_PARAMS)
|
|
self.units.append(unit)
|
|
|
|
def train(self, vc: Mat, status: Status = None):
|
|
for unit in self.units:
|
|
train(unit, vc, status)
|
|
# Update context portion of vc
|
|
_c = unit.forward(vc)
|
|
vc[1:, WIN*vocab_size():] = _c[0:-1,:]
|
|
|
|
def forward_step(self, _v: Mat, _c: Mat) -> tuple[Mat, Mat]:
|
|
_vc = concat(_v.reshape([1, WIN*vocab_size()]), _c.reshape([1, H_SIZE]), axis=1)
|
|
for unit in self.units:
|
|
_c = unit.forward(_vc)
|
|
_vc = unit.reconstruct(_c)
|
|
return split(_vc.flatten(), H_SIZE)
|
|
|
|
def forward(self, x: Mat):
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if 1:
|
|
# Test of helper function
|
|
# ToDo: move to tests
|
|
vec = idx2vec(ch2idx("JENS"))
|
|
vec = clamp(vec, axis=1)
|
|
indices = vec2idx(vec, axis=1)
|
|
ch_str = idx2ch(indices)
|
|
print(ch_str)
|
|
|
|
win_text = to_window(TEXT)
|
|
print(win_text)
|
|
|
|
batch = to_batch(win_text)
|
|
print(batch)
|
|
|
|
|
|
model = RnnModel(name='JayRnn', work_dir='results')
|
|
model.init(0.01)
|
|
model.load()
|
|
|
|
# vc contains vis + context
|
|
# context will be updated after training
|
|
c_train = np.zeros([len(batch), H_SIZE])
|
|
vc_train = concat(batch, c_train, axis=1)
|
|
for k in range(NUM_ITERATIONS):
|
|
model.train(vc_train, status=Status())
|
|
model.save()
|
|
|
|
# test the model
|
|
seed = str2vec('HA^')
|
|
v_forward = seed.flatten()
|
|
context = c_train[0]
|
|
for i in range(len(TEXT)+5):
|
|
v_mat = v_forward.reshape([WIN, vocab_size()])
|
|
print(f"Forward {i:02d}: {vec2str(v_mat, axis=1)}")
|
|
v_predict, context = model.forward_step(v_mat, context)
|
|
v_predict_clamp = clamp(v_predict.reshape([WIN, vocab_size()]), axis=1)
|
|
print(f"Predict {i:02d}: {vec2str(v_predict_clamp, axis=1)}")
|
|
v_forward = shift_left(v_predict_clamp.reshape([1, WIN*vocab_size()]), vocab_size())
|
|
|