getting Rnn to work

src/stack/rnn_helper.py

@@ -1,33 +1,43 @@
 from rbm.matrix import Mat, np
 
-VOCAB = ' .!?ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
+VOCAB = '^ .!?ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
 _CH2IDX = {ch: i for i, ch in enumerate(VOCAB)}
 
-def vec2idx(vec: Mat) -> int:
-	index = np.argmax(vec)
-	return int(index)
+def clamp(vec: Mat, axis=0):
+	indices = vec2idx(vec, axis=axis)
+	return idx2vec(indices)
 
-def idx2vec(index: int) -> np.ndarray:
-	result = np.zeros([1, len(VOCAB)])
-	result[:, index] = 1
+def vec2idx(vec: Mat, axis=0) -> Mat:
+	indices = np.argmax(vec, axis=axis)
+	return indices
+
+def idx2vec(indices: Mat) -> Mat:
+	result = np.zeros([len(indices), len(VOCAB)])
+	for i, index in enumerate(indices):
+		result[i, int(index)] = 1
 	return result
 
-def ch2idx(ch: str) -> int:
-	idx = _CH2IDX[ch]
+def ch2idx(ch_str: str) -> Mat:
+	idx = np.zeros(len(ch_str))
+	for i, ch in enumerate(ch_str):
+		idx[i] = int(_CH2IDX[ch])
 	return idx
 
-def idx2ch(idx: int) -> str:
-	return VOCAB[idx]
+def idx2ch(indices: Mat) -> str:
+	result = ''
+	for index in indices:
+		result += VOCAB[int(index)]
+	return result
 
 def vocab_size():
 	return len(VOCAB)
 
-def concat(v: np.ndarray, c: np.ndarray) -> np.ndarray:
-	return np.concatenate([v, c], axis=1)
+def concat(v: np.ndarray, c: np.ndarray, axis=0) -> np.ndarray:
+	return np.concatenate([v, c], axis=axis)
 
-def split(vc: np.ndarray, v_size: np.ndarray|int) -> tuple[np.ndarray, np.ndarray]:
-	n = len(v_size) if isinstance(v_size, np.ndarray) else v_size
-	return vc[:n], vc[n:]
+def split(vc: np.ndarray, h_size: int, axis=0) -> tuple[np.ndarray, np.ndarray]:
+	v_len = vc.shape[axis]-h_size
+	return vc[0:v_len], vc[v_len:]
 
 def shift_right(m: np.ndarray, amount: int = 1) -> np.ndarray:
 	return np.hstack([np.zeros((m.shape[0], amount)), m[:, :-amount]])