numpy_lstm.py

'''
basic LSTM written by claude
'''


import numpy as np
from tokenizer import make_train_and_test_shakepeare

class LSTMCell:
    """A single LSTM cell implementation from scratch."""
    
    def __init__(self, input_size, hidden_size):
        """
        Initialize LSTM cell parameters.
        
        Args:
            input_size: Dimension of input vector
            hidden_size: Dimension of hidden state and cell state
        """
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        # Initialize weights (Xavier initialization)
        scale = 1.0 / np.sqrt(input_size + hidden_size)
        
        # Forget gate weights
        self.Wf = np.random.randn(hidden_size, input_size + hidden_size) * scale
        self.bf = np.zeros((hidden_size, 1))
        
        # Input gate weights
        self.Wi = np.random.randn(hidden_size, input_size + hidden_size) * scale
        self.bi = np.zeros((hidden_size, 1))
        
        # Candidate cell state weights
        self.Wc = np.random.randn(hidden_size, input_size + hidden_size) * scale
        self.bc = np.zeros((hidden_size, 1))
        
        # Output gate weights
        self.Wo = np.random.randn(hidden_size, input_size + hidden_size) * scale
        self.bo = np.zeros((hidden_size, 1))
        
        # Gradients (initialized to zeros)
        self.dWf = np.zeros_like(self.Wf)
        self.dbf = np.zeros_like(self.bf)
        self.dWi = np.zeros_like(self.Wi)
        self.dbi = np.zeros_like(self.bi)
        self.dWc = np.zeros_like(self.Wc)
        self.dbc = np.zeros_like(self.bc)
        self.dWo = np.zeros_like(self.Wo)
        self.dbo = np.zeros_like(self.bo)
    
    def sigmoid(self, x):
        """Sigmoid activation function."""
        return 1 / (1 + np.exp(-np.clip(x, -500, 500)))
    
    def tanh(self, x):
        """Tanh activation function."""
        return np.tanh(x)
    
    def forward(self, x, h_prev, c_prev):
        """
        Forward pass through LSTM cell.
        
        Args:
            x: Input at current time step (input_size, 1)
            h_prev: Previous hidden state (hidden_size, 1)
            c_prev: Previous cell state (hidden_size, 1)
            
        Returns:
            h: Current hidden state
            c: Current cell state
            cache: Values needed for backward pass
        """
        # Concatenate previous hidden state and current input
        concat = np.vstack((h_prev, x))
        
        # Forget gate
        f = self.sigmoid(np.dot(self.Wf, concat) + self.bf)
        
        # Input gate
        i = self.sigmoid(np.dot(self.Wi, concat) + self.bi)
        
        # Candidate cell state
        c_tilde = self.tanh(np.dot(self.Wc, concat) + self.bc)
        
        # Update cell state
        c = f * c_prev + i * c_tilde
        
        # Output gate
        o = self.sigmoid(np.dot(self.Wo, concat) + self.bo)
        
        # Update hidden state
        h = o * self.tanh(c)
        
        # Cache values for backward pass
        cache = (x, h_prev, c_prev, concat, f, i, c_tilde, c, o)
        
        return h, c, cache
    
    def backward(self, dh, dc_next, cache):
        """
        Backward pass through LSTM cell.
        
        Args:
            dh: Gradient of loss w.r.t. hidden state (hidden_size, 1)
            dc_next: Gradient of loss w.r.t. cell state from next time step (hidden_size, 1)
            cache: Values cached during forward pass
            
        Returns:
            dx: Gradient w.r.t. input
            dh_prev: Gradient w.r.t. previous hidden state
            dc_prev: Gradient w.r.t. previous cell state
        """
        # Unpack cache
        x, h_prev, c_prev, concat, f, i, c_tilde, c, o = cache
        
        # Gradient from output gate
        dc = dh * o * (1 - self.tanh(c)**2) + dc_next
        
        # Gradient of output gate
        do = dh * self.tanh(c)
        do_input = do * o * (1 - o)  # sigmoid derivative
        
        # Gradient of cell state components
        dc_tilde = dc * i
        dc_tilde_input = dc_tilde * (1 - c_tilde**2)  # tanh derivative
        
        di = dc * c_tilde
        di_input = di * i * (1 - i)  # sigmoid derivative
        
        df = dc * c_prev
        df_input = df * f * (1 - f)  # sigmoid derivative
        
        # Gradients w.r.t. concatenated input [h_prev, x]
        dconcat_f = np.dot(self.Wf.T, df_input)
        dconcat_i = np.dot(self.Wi.T, di_input)
        dconcat_c = np.dot(self.Wc.T, dc_tilde_input)
        dconcat_o = np.dot(self.Wo.T, do_input)
        dconcat = dconcat_f + dconcat_i + dconcat_c + dconcat_o
        
        # Split gradient for h_prev and x
        dh_prev = dconcat[:self.hidden_size, :]
        dx = dconcat[self.hidden_size:, :]
        
        # Gradient w.r.t. previous cell state
        dc_prev = dc * f
        
        # Accumulate gradients for weights and biases
        self.dWf += np.dot(df_input, concat.T)
        self.dbf += df_input
        
        self.dWi += np.dot(di_input, concat.T)
        self.dbi += di_input
        
        self.dWc += np.dot(dc_tilde_input, concat.T)
        self.dbc += dc_tilde_input
        
        self.dWo += np.dot(do_input, concat.T)
        self.dbo += do_input
        
        return dx, dh_prev, dc_prev
    
    def reset_gradients(self):
        """Reset accumulated gradients to zero."""
        self.dWf.fill(0)
        self.dbf.fill(0)
        self.dWi.fill(0)
        self.dbi.fill(0)
        self.dWc.fill(0)
        self.dbc.fill(0)
        self.dWo.fill(0)
        self.dbo.fill(0)
    
    def update_weights(self, learning_rate):
        """Update weights using accumulated gradients."""
        self.Wf -= learning_rate * self.dWf
        self.bf -= learning_rate * self.dbf
        
        self.Wi -= learning_rate * self.dWi
        self.bi -= learning_rate * self.dbi
        
        self.Wc -= learning_rate * self.dWc
        self.bc -= learning_rate * self.dbc
        
        self.Wo -= learning_rate * self.dWo
        self.bo -= learning_rate * self.dbo


class LSTM:
    """Multi-layer LSTM network."""
    
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        """
        Initialize LSTM network.
        
        Args:
            input_size: Dimension of input vector
            hidden_size: Dimension of hidden states
            output_size: Dimension of output vector
            num_layers: Number of LSTM layers
        """
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.num_layers = num_layers
        
        # Create LSTM cells
        self.cells = []
        for layer in range(num_layers):
            if layer == 0:
                cell = LSTMCell(input_size, hidden_size)
            else:
                cell = LSTMCell(hidden_size, hidden_size)
            self.cells.append(cell)
        
        # Output layer weights
        self.Wy = np.random.randn(output_size, hidden_size) * 0.01
        self.by = np.zeros((output_size, 1))
        
        # Output layer gradients
        self.dWy = np.zeros_like(self.Wy)
        self.dby = np.zeros_like(self.by)
    
    def forward(self, inputs):
        """
        Forward pass through LSTM network.
        
        Args:
            inputs: List of input vectors, one for each time step
                    Each input shape: (input_size, 1)
                    
        Returns:
            outputs: List of output vectors for each time step
            caches: Cached values for backward pass
        """
        T = len(inputs)
        
        # Initialize hidden and cell states for each layer
        h = [np.zeros((self.hidden_size, 1)) for _ in range(self.num_layers)]
        c = [np.zeros((self.hidden_size, 1)) for _ in range(self.num_layers)]
        
        # Store all hidden and cell states for backward pass
        all_h = [[h[layer].copy() for layer in range(self.num_layers)] for _ in range(T + 1)]
        all_c = [[c[layer].copy() for layer in range(self.num_layers)] for _ in range(T + 1)]
        
        outputs = []
        caches = []
        
        # Process each time step
        for t in range(T):
            layer_input = inputs[t]
            layer_caches = []
            
            # Forward through each LSTM layer
            for layer in range(self.num_layers):
                h[layer], c[layer], cache = self.cells[layer].forward(
                    layer_input, h[layer], c[layer]
                )
                layer_caches.append(cache)
                layer_input = h[layer]  # Output of this layer is input to next
            
            # Store states
            for layer in range(self.num_layers):
                all_h[t + 1][layer] = h[layer].copy()
                all_c[t + 1][layer] = c[layer].copy()
            
            # Output layer
            y = np.dot(self.Wy, h[-1]) + self.by
            
            outputs.append(y)
            caches.append(layer_caches)
        
        return outputs, caches, all_h, all_c
    
    def backward(self, doutputs, caches, all_h, all_c):
        """
        Backward pass through LSTM network (BPTT).
        
        Args:
            doutputs: List of gradients w.r.t. outputs for each time step
            caches: Cached values from forward pass
            all_h: All hidden states from forward pass
            all_c: All cell states from forward pass
            
        Returns:
            dinputs: List of gradients w.r.t. inputs for each time step
        """
        T = len(doutputs)
        
        # Reset gradients
        for cell in self.cells:
            cell.reset_gradients()
        self.dWy.fill(0)
        self.dby.fill(0)
        
        # Initialize gradients for hidden and cell states
        dh_next = [np.zeros((self.hidden_size, 1)) for _ in range(self.num_layers)]
        dc_next = [np.zeros((self.hidden_size, 1)) for _ in range(self.num_layers)]
        
        dinputs = []
        
        # Backpropagate through time
        for t in reversed(range(T)):
            # Gradient from output layer
            dy = doutputs[t]
            self.dWy += np.dot(dy, all_h[t + 1][-1].T)
            self.dby += dy
            
            dh_next[-1] += np.dot(self.Wy.T, dy)
            
            # Backpropagate through LSTM layers (in reverse order)
            layer_caches = caches[t]
            dx = None
            
            for layer in reversed(range(self.num_layers)):
                cache = layer_caches[layer]
                dx, dh_prev, dc_prev = self.cells[layer].backward(
                    dh_next[layer], dc_next[layer], cache
                )
                
                # Update gradients for next iteration
                dh_next[layer] = dh_prev
                dc_next[layer] = dc_prev
                
                # If not the first layer, pass gradient to previous layer
                if layer > 0:
                    dh_next[layer - 1] += dx
            
            dinputs.insert(0, dx)
        
        return dinputs
    
    def update_weights(self, learning_rate):
        """Update all weights in the network."""
        for cell in self.cells:
            cell.update_weights(learning_rate)
        
        self.Wy -= learning_rate * self.dWy
        self.by -= learning_rate * self.dby
    
    def train_step(self, inputs, targets, learning_rate=0.01):
        """
        Single training step.
        
        Args:
            inputs: List of input vectors
            targets: List of target vectors
            learning_rate: Learning rate for weight updates
            
        Returns:
            loss: Mean squared error loss
        """
        # Forward pass
        outputs, caches, all_h, all_c = self.forward(inputs)
        
        # Compute loss (MSE) and gradients
        loss = 0
        doutputs = []
        for t in range(len(outputs)):
            error = outputs[t] - targets[t]
            loss += np.sum(error ** 2)
            doutputs.append(2 * error)  # Gradient of MSE
        
        loss /= len(outputs)
        
        # Backward pass
        self.backward(doutputs, caches, all_h, all_c)
        
        # Update weights
        self.update_weights(learning_rate)
        
        return loss


# Example usage with training
if __name__ == "__main__":
    # Create a simple LSTM
    input_size = 5
    hidden_size = 10
    output_size = 1
    
    lstm = LSTM(input_size, hidden_size, output_size, num_layers=2)
    
    # Create sample training data
    sequence_length = 8
    num_sequences = 5
    
    train_tok_li, train_ids_li, test_tok_li, test_ids_li = make_train_and_test_shakepeare()
    train_inputs = []
    train_targets = []
    test_inputs = []
    test_targets = []
    num_sequences = len(train_tok_li)
    for seq in train_ids_li:
        train_inputs.append(seq[:-1]) # TODO try different output lengths
        train_targets.append(seq[-1:])
    for seq in test_ids_li:
        test_inputs.append(seq[:-1])
        test_targets.append(seq[-1:])

    print("Training LSTM...")
    for epoch in range(100):
        total_loss = 0
        
        for inputs, targets in zip(train_inputs,train_targets):
            # Train
            loss = lstm.train_step(inputs, targets, learning_rate=0.001)
            total_loss += loss
        
        if epoch % 20 == 0:
            avg_loss = total_loss / num_sequences
            print(f"Epoch {epoch}, Average Loss: {avg_loss:.4f}")
    
    print("\nTraining complete!")
    
    # Test forward pass
    outputs, _, _, _ = lstm.forward(test_inputs)
    
    print(f"\nTest sequence length: {len(test_inputs)}")
    print(f"Output sequence length: {len(outputs)}")
    print(f"First output:\n{outputs[0]}")