Add RNN model with attention

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1 @@
+{"model_type": "RNN", "vocab_size": 250002, "hidden_size": 256, "output_size": 2, "cell_type": "RNN", "architecture": "SimpleRecurrentNetworkWithAttention"}

modeling.py

@@ -0,0 +1,301 @@
+
+import time
+from datasets import load_dataset
+import numpy as np
+import transformers
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+    
+class custom_RNNCell(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, device='cpu'):
+        
+        # Initialize a basic RNN cell with Xavier-initialized weights.
+        # :param input_size: Number of input features.
+        # :param hidden_size: Number of units in the hidden layer.
+        #:param device: Device to place the tensors on.
+        
+        super(custom_RNNCell, self).__init__()
+        self.hidden_size = hidden_size
+        self.device = device
+
+        # Xavier initialization limits
+        fan_in_Wx = input_size
+        fan_out_Wx = hidden_size
+        limit_Wx = np.sqrt(6 / (fan_in_Wx + fan_out_Wx))
+
+        fan_in_Wh = hidden_size
+        fan_out_Wh = hidden_size
+        limit_Wh = np.sqrt(6 / (fan_in_Wh + fan_out_Wh))
+
+        # Convert weights to PyTorch Parameters
+        self.Wx = nn.Parameter(torch.empty(input_size, hidden_size, device=device))
+        self.Wh = nn.Parameter(torch.empty(hidden_size, hidden_size, device=device))
+        self.bh = nn.Parameter(torch.zeros(hidden_size, device=device))
+
+        # Initialize using Xavier uniform
+        nn.init.uniform_(self.Wx, -limit_Wx, limit_Wx)
+        nn.init.uniform_(self.Wh, -limit_Wh, limit_Wh)
+
+    def forward(self, input_t: torch.Tensor, h_prev: torch.Tensor) -> torch.Tensor:
+        
+        # Forward pass for a basic RNN cell.
+        # :param input_t: Input at time step t (batch_size x input_size).
+        # :param h_prev: Hidden state from previous time step (batch_size x hidden_size).
+        # :return: Updated hidden state.
+        
+        h_t = torch.tanh(torch.mm(input_t, self.Wx) + torch.mm(h_prev, self.Wh) + self.bh)
+        return h_t
+
+class custom_GRUCell(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, device='cpu'):
+        
+        # Initialize a GRU cell with Xavier-initialized weights.
+        # :param input_size: Number of input features.
+        # :param hidden_size: Number of units in the hidden layer.
+        # :param device: The device to run the computations on.
+        
+        super(custom_GRUCell, self).__init__()
+        self.hidden_size = hidden_size
+        self.device = device
+
+        # Xavier initialization limits
+        fan_in = input_size + hidden_size
+        fan_out = hidden_size
+        limit = (6 / (fan_in + fan_out)) ** 0.5
+
+        # Weight matrices for update gate, reset gate, and candidate hidden state
+        self.Wz = nn.Parameter(torch.empty(input_size + hidden_size, hidden_size, device=device))  # Update gate
+        self.Wr = nn.Parameter(torch.empty(input_size + hidden_size, hidden_size, device=device))  # Reset gate
+        self.Wh = nn.Parameter(torch.empty(input_size + hidden_size, hidden_size, device=device))  # Candidate hidden state
+
+        # Apply Xavier initialization
+        nn.init.uniform_(self.Wz, -limit, limit)
+        nn.init.uniform_(self.Wr, -limit, limit)
+        nn.init.uniform_(self.Wh, -limit, limit)
+
+        # Biases for each gate, initialized to zeros
+        self.bz = nn.Parameter(torch.zeros(hidden_size, device=device))
+        self.br = nn.Parameter(torch.zeros(hidden_size, device=device))
+        self.bh = nn.Parameter(torch.zeros(hidden_size, device=device))
+
+    def forward(self, input_t: torch.Tensor, h_prev: torch.Tensor) -> torch.Tensor:
+        
+        # Forward pass for a single GRU cell.
+        # :param input_t: Input at time step t (batch_size x input_size).
+        # :param h_prev: Hidden state from the previous time step (batch_size x hidden_size).
+        # :return: Updated hidden state.
+        
+        # Concatenate input and previous hidden state
+        concat = torch.cat((input_t, h_prev), dim=1)
+
+        # Update gate
+        z_t = torch.sigmoid(torch.matmul(concat, self.Wz) + self.bz)
+
+        # Reset gate
+        r_t = torch.sigmoid(torch.matmul(concat, self.Wr) + self.br)
+
+        # Candidate hidden state
+        concat_reset = torch.cat((input_t, r_t * h_prev), dim=1)
+        h_hat_t = torch.tanh(torch.matmul(concat_reset, self.Wh) + self.bh)
+
+        # Compute final hidden state
+        h_t = (1 - z_t) * h_prev + z_t * h_hat_t
+
+        return h_t
+
+class custom_LSTMCell(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, device='cpu'):
+        super(custom_LSTMCell, self).__init__()
+        self.hidden_size = hidden_size
+        self.device = device
+
+        # Initialize LSTM weights and biases using Xavier initialization
+        self.Wf = nn.Parameter(torch.empty(input_size + hidden_size, hidden_size, device=device)) # Forget gate (W_f)
+        self.Wi = nn.Parameter(torch.empty(input_size + hidden_size, hidden_size, device=device)) # Input gate (W_i)
+        self.Wc = nn.Parameter(torch.empty(input_size + hidden_size, hidden_size, device=device)) # Candidate cell state (W_c~)
+        self.Wo = nn.Parameter(torch.empty(input_size + hidden_size, hidden_size, device=device)) # Output gate (W_o)
+
+        # Apply Xavier initialization
+        nn.init.xavier_uniform_(self.Wf)
+        nn.init.xavier_uniform_(self.Wi)
+        nn.init.xavier_uniform_(self.Wc)
+        nn.init.xavier_uniform_(self.Wo)
+
+        # Initialize biases
+        self.bf = nn.Parameter(torch.zeros(hidden_size, device=device)) # Forget gate bias (b_f)
+        self.bi = nn.Parameter(torch.zeros(hidden_size, device=device)) # Input gate bias (b_i)
+        self.bc = nn.Parameter(torch.zeros(hidden_size, device=device)) # Candidate state bias (b_c~)
+        self.bo = nn.Parameter(torch.zeros(hidden_size, device=device)) # Output gate bias (b_o)
+
+        # Initialize forget gate bias to positive value to help with training
+        nn.init.constant_(self.bf, 1.0)
+
+    def forward(self, x_t: torch.Tensor, h_prev: torch.Tensor, c_prev: torch.Tensor) -> tuple:
+        
+        # Forward pass for a single LSTM cell.
+        # :param x_t: Input at the current time step (batch_size x input_size).
+        # :param h_prev: Previous hidden state (batch_size x hidden_size).
+        # :param c_prev: Previous cell state (batch_size x hidden_size).
+        # :return: Tuple of new hidden state (h_t) and new cell state (c_t).
+        
+        # Concatenate input and previous hidden state (x_t and h_{t-1})
+        concat = torch.cat((x_t, h_prev), dim=1)
+
+        # Forget gate: decides what to remove from the cell state
+        f_t = torch.sigmoid(torch.matmul(concat, self.Wf) + self.bf) # Forget gate (σ) -> f_t
+
+        # Input gate: decides what to add to the cell state
+        i_t = torch.sigmoid(torch.matmul(concat, self.Wi) + self.bi) # Input gate (σ) -> i_t
+
+        # Candidate cell state: new information to potentially add to the cell state
+        c_hat_t = torch.tanh(torch.matmul(concat, self.Wc) + self.bc) # Candidate cell state (tanh) -> c~_t
+
+        # Update cell state: new cell state (c_t) based on previous state and gates
+        c_t = f_t * c_prev + i_t * c_hat_t # Cell state update -> c_t
+
+        # Output gate: decides what the next hidden state should be
+        o_t = torch.sigmoid(torch.matmul(concat, self.Wo) + self.bo) # Output gate (σ) -> o_t
+
+        # New hidden state (h_t): based on cell state and output gate
+        h_t = o_t * torch.tanh(c_t) # Hidden state update -> h_t
+
+        # Return the new hidden state (h_t) and cell state (c_t)
+        return h_t, c_t
+
+class RecurrentLayer(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, cell_type: str = 'RNN', device='cpu'):
+        super(RecurrentLayer, self).__init__()
+        self.hidden_size = hidden_size
+        self.device = device
+        self.cell_type = cell_type
+
+        # Initialize the appropriate cell type
+        if cell_type == 'RNN':
+            self.cell = nn.RNN(input_size, hidden_size, batch_first=True, bidirectional=True)
+        elif cell_type == 'custom_RNN':
+            self.cell = custom_RNNCell(input_size, hidden_size)
+        elif cell_type == 'GRU':
+            self.cell = nn.GRU(input_size, hidden_size, batch_first=True, bidirectional=True)
+        elif cell_type == 'custom_GRU':
+            self.cell = custom_GRUCell(input_size, hidden_size, device)
+        elif cell_type == 'LSTM':
+            self.cell = nn.LSTMCell(input_size, hidden_size)
+        elif cell_type == 'custom_LSTM':
+            self.cell = custom_LSTMCell(input_size, hidden_size, device)
+        else:
+            raise ValueError("Unsupported cell type")
+
+    def forward(self, inputs: torch.Tensor) -> tuple:
+        
+        # Forward pass through the recurrent layer for a sequence of inputs.
+        # Returns a tuple of (output, last_hidden_state) to match PyTorch's interface.
+        
+        batch_size, seq_len, _ = inputs.shape
+        
+        # Initialize hidden states
+        h_forward = torch.zeros(batch_size, self.hidden_size, device=self.device)
+        h_backward = torch.zeros(batch_size, self.hidden_size, device=self.device)
+        
+        if self.cell_type == 'custom_LSTM':
+            c_forward = torch.zeros(batch_size, self.hidden_size, device=self.device)
+            c_backward = torch.zeros(batch_size, self.hidden_size, device=self.device)
+        
+        # Lists to store outputs for both directions
+        forward_outputs = []
+        backward_outputs = []
+        
+        # Forward pass
+        h = h_forward
+        c = c_forward if self.cell_type == 'custom_LSTM' else None
+        for t in range(seq_len):
+            if self.cell_type == 'custom_LSTM':
+                h, c = self.cell(inputs[:, t], h, c)
+            else:
+                h = self.cell(inputs[:, t], h)
+            forward_outputs.append(h)
+            
+        # Backward pass
+        h = h_backward
+        c = c_backward if self.cell_type == 'custom_LSTM' else None
+        for t in range(seq_len - 1, -1, -1):
+            if self.cell_type == 'custom_LSTM':
+                h, c = self.cell(inputs[:, t], h, c)
+            else:
+                h = self.cell(inputs[:, t], h)
+            backward_outputs.insert(0, h)
+            
+        # Stack and concatenate outputs
+        forward_output = torch.stack(forward_outputs, dim=1)  # [batch_size, seq_len, hidden_size]
+        backward_output = torch.stack(backward_outputs, dim=1)  # [batch_size, seq_len, hidden_size]
+        output = torch.cat((forward_output, backward_output), dim=2)  # [batch_size, seq_len, 2*hidden_size]
+        
+        # Create final hidden state (concatenated forward and backward)
+        final_hidden = torch.stack([forward_outputs[-1], backward_outputs[-1]], dim=0)  # [2, batch_size, hidden_size]
+        
+        return output, final_hidden
+    
+class Attention(nn.Module):
+    def __init__(self, hidden_size):
+        super(Attention, self).__init__()
+        self.W1 = nn.Linear(hidden_size, hidden_size)
+        self.W2 = nn.Linear(hidden_size, hidden_size)
+        self.v = nn.Linear(hidden_size, 1, bias=False)
+
+    def forward(self, hidden, encoder_outputs):
+        # hidden: [batch_size, hidden_size]
+        # encoder_outputs: [batch_size, sequence_len, hidden_size]
+        sequence_len = encoder_outputs.shape[1]
+        hidden = hidden.unsqueeze(1).repeat(1, sequence_len, 1)
+
+        energy = torch.tanh(self.W1(encoder_outputs) + self.W2(hidden))
+        attention = self.v(energy).squeeze(2)  # [batch_size, sequence_len]
+        attention_weights = torch.softmax(attention, dim=1)
+
+        # Apply attention weights to encoder outputs to get context vector
+        context = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs).squeeze(1)
+        return context, attention_weights
+    
+class SimpleRecurrentNetworkWithAttention(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size, cell_type='RNN'):
+        super(SimpleRecurrentNetworkWithAttention, self).__init__()
+
+        self.embedding = nn.Embedding(input_size, hidden_size)
+        self.attention = Attention(hidden_size * 2)  # Use hidden_size * 2 for bidirectional LSTM
+        self.cell_type = cell_type
+
+        if cell_type == 'RNN':
+            self.cell = nn.RNN(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        elif cell_type == 'custom_RNN':
+            self.cell = RecurrentLayer(hidden_size, hidden_size, cell_type="custom_RNN") #custom_RNNCell(input_size, hidden_size)
+        elif cell_type == 'GRU':
+            self.cell = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        elif cell_type == 'custom_GRU':
+            self.cell = RecurrentLayer(hidden_size, hidden_size, cell_type="custom_GRU")
+        elif cell_type == 'LSTM':
+            self.cell = nn.LSTM(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        elif cell_type == 'custom_LSTM':
+            self.cell = RecurrentLayer(hidden_size, hidden_size, cell_type="custom_LSTM")
+        else:
+            raise ValueError("Unsupported cell type. Choose from 'RNN', 'custom_RNN', 'GRU', 'custom_GRU', 'LSTM' or 'custom_LSTM'.")
+
+        self.fc = nn.Linear(hidden_size * 2, output_size)  # hidden_size * 2 for bidirectional
+
+    def forward(self, inputs):
+        embedded = self.embedding(inputs)
+        rnn_output, hidden = self.cell(embedded)
+
+        if isinstance(hidden, tuple):  # LSTM returns (hidden, cell_state)
+            hidden = hidden[0]
+
+        # Since it's bidirectional, get the last layer's forward and backward hidden states
+        hidden = torch.cat((hidden[-2], hidden[-1]), dim=1)  # Concatenate forward and backward hidden states
+
+        # Apply attention to the concatenated hidden state
+        context, attention_weights = self.attention(hidden, rnn_output)
+
+        # Pass the context vector to the fully connected layer
+        output = self.fc(context)
+
+        return output, attention_weights

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d75a21bee6da43507a4799f32ffc505bccaca9c1d9f514341ec53c5d2a95e9e
+size 259167710

special_tokens_map.json

@@ -0,0 +1,15 @@
+{
+  "bos_token": "<s>",
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "unk_token": "<unk>"
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:affcfb1f45c4b14a70a6589c3d153b430ed4309e5a6613a88dab64d5a923a5d6
+size 17082925

tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": {
+    "__type": "AddedToken",
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "XLMRobertaTokenizer",
+  "unk_token": "<unk>"
+}