model.py

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import PreTrainedModel

from typing import List

from .config import LidirlLSTMConfig

def torch_max_no_pads(model_out, lengths):
    indices = torch.arange(model_out.size(1)).to(model_out.device)
    mask = (indices < lengths.view(-1, 1)).unsqueeze(-1).expand(model_out.size())
    model_out = torch.where(mask, model_out, torch.tensor(-1e9))
    max_pool = torch.max(model_out, 1)[0]
    return max_pool


class ProjectionLayer(nn.Module):
    """
        Noise-aware labels layer or traditional linear projection
    """

    def __init__(self,
        hidden_dim : int,
        label_size : int,
        montecarlo_layer : bool = False):
        super().__init__()
        self.montecarlo_layer = montecarlo_layer
        if montecarlo_layer:
            self.proj = MCSoftmaxDenseFA(hidden_dim, label_size, 1, logits_only=True)
        else:
            # eventually I'm considering making this a variable size classifier so it's a sequence of layers even though it's just one
            self.projs = [
                nn.Linear(hidden_dim, label_size)
            ]
            self.proj = nn.Sequential(*self.projs)
        self.init_layer()

    def forward(self, x):
        return self.proj(x)

    def init_layer(self, pi : float = 0.01):
        """
            Initialize the final classification layer so all predictions are close to 0
        """

        if self.montecarlo_layer:
            self.proj.init_layer(pi)
            return
        
        for layer in self.proj.modules():
            if isinstance(layer, nn.Linear):
                nn.init.normal_(layer.weight, mean=0.0, std=0.01)
                if layer.bias is not None:
                    # Bias: -log((1 - pi) / pi)
                    bias_value = -math.log((1 - pi) / pi)
                    nn.init.constant_(layer.bias, bias_value)


class MinLSTMCell(nn.Module):
    """
        https://arxiv.org/pdf/2410.01201
        https://github.com/YecanLee/min-LSTM-torch/blob/main/minLSTMcell.py
    bidirectional and parallel
    hold layer depth and sweep out the other dimensions
    """
    def __init__(self, 
                 embed_dim,
                 hidden_dim):
        super(MinLSTMCell, self).__init__()
        self.embed_dim = embed_dim
        self.hidden_dim = hidden_dim
        self.output_dim = embed_dim
        
        # Initialize the linear layers for the forget gate, input gate, and hidden state transformation
        self.linear_f = nn.Linear(embed_dim, hidden_dim)
        self.linear_i = nn.Linear(embed_dim, hidden_dim)
        self.linear_h = nn.Linear(embed_dim, hidden_dim)

    def parallel_scan_log(self, log_coeffs, log_values):
        # log_coeffs: (batch_size, seq_len, input_size)
        # log_values: (batch_size, seq_len + 1, input_size)
        a_star = F.pad(torch.cumsum(log_coeffs, dim=1), (0, 0, 1, 0))
        log_h0_plus_b_star = torch.logcumsumexp(
            log_values - a_star, dim=1)
        log_h = a_star + log_h0_plus_b_star
        return torch.exp(log_h)[:, 1:]

    def g(self, x):
        return torch.where(x >= 0, x+0.5, torch.sigmoid(x))

    def log_g(self, x):
        return torch.where(x >= 0, (F.relu(x)+0.5).log(), -F.softplus(-x))

    def forward(self, inputs):
        h_init = torch.zeros(inputs.size(0), 1, self.hidden_dim, device=inputs.device)

        diff = F.softplus(-self.linear_f(inputs)) - F.softplus(-self.linear_i(inputs))

        log_f = -F.softplus(diff)
        log_i = -F.softplus(-diff)
        log_h_0 = torch.log(h_init)

        log_tilde_h = self.log_g(self.linear_h(inputs))

        h = self.parallel_scan_log(log_f, torch.cat([log_h_0, log_i + log_tilde_h], dim=1))
        return h


class LSTMBlock(nn.Module):
    def __init__(self,
                    embed_dim : int = 512,
                    hidden_dim : int = 2048,
                    num_layers : int = 6,
                    dropout : float = 0.1,
                    bidirectional : bool = False
                 ):
        super(LSTMBlock, self).__init__()

        self.layers = []
        last_dim = embed_dim
        for _ in range(num_layers):
            self.layers.append(MinLSTMCell(last_dim, hidden_dim))
            self.layers.append(nn.LayerNorm(hidden_dim, elementwise_affine=True))
            self.layers.append(nn.GELU())
            self.layers.append(nn.Dropout(dropout))
            last_dim = hidden_dim
        self.model = nn.Sequential(*self.layers)
        self.bidirectionality_term = 2 if bidirectional else 1
        self.output_dim = hidden_dim * self.bidirectionality_term
        self.bidirectional = bidirectional

    def flip_sequence(self, inputs, lengths):
        # Here we want to flip the sequence but keep the right-padding
        # We can do this by flipping the sequence and then flipping the padding
        new = []
        for inp, leng in zip(inputs, lengths):
            new.append(inp[:leng].flip(0))
        return pad_sequence(new, batch_first=True).to(inputs.device)

    def forward(self, inputs, lengths):
        encoding = self.model(inputs)
        last_token = encoding[torch.arange(encoding.size(0)), lengths - 1].view(inputs.size(0), 1, -1)
        if self.bidirectional:
            reverse_sequence = self.flip_sequence(inputs, lengths)
            reverse_encoding = self.model(reverse_sequence)
            reverse_last_token = reverse_encoding[torch.arange(reverse_encoding.size(0)), lengths - 1].view(inputs.size(0), 1, -1)
            last_token = torch.cat((last_token, reverse_last_token), dim=-1)

        return last_token, torch.ones((inputs.size(0), 1), device=inputs.device, dtype=torch.long)


class LidirlLSTM(PreTrainedModel):
    """
        Defines the Lidirl LSTM Model
    """

    config_class = LidirlLSTMConfig
    def __init__(self, config):
        super().__init__(config)
        
        self.encoder = LSTMBlock(
            embed_dim = config.embed_dim,
            hidden_dim = config.hidden_dim,
            num_layers = config.num_layers,
            dropout = config.dropout,
            bidirectional = config.bidirectional
        )
        self.embed_layer = nn.Embedding(config.vocab_size, config.embed_dim)
        self.proj = ProjectionLayer(self.encoder.output_dim, config.label_size, config.montecarlo_layer)

        self.label_size = config.label_size
        self.max_length = config.max_length
        self.multilabel = config.multilabel
        self.monte_carlo = config.montecarlo_layer

        self.labels = ["" for _ in config.labels]
        for key, value in config.labels.items():
            self.labels[value] = key

    def forward(self, inputs, lengths):
        inputs = inputs[:, :self.max_length]
        lengths = lengths.clamp(max=self.max_length)

        embeddings = self.embed_layer(inputs)
        encoding, lengths = self.encoder(embeddings, lengths=lengths)
        max_pool = torch_max_no_pads(encoding, lengths)
        projection = self.proj(max_pool)

        return projection

    def __call__(self, inputs, lengths):
        # this is inference only model
        with torch.no_grad():
            logits = self.forward(inputs, lengths)
            if self.multilabel:
                probs = torch.sigmoid(logits)
            else:
                probs = torch.softmax(logits, dim=-1)
        return probs

    def predict(self, inputs, lengths, threshold=0.5, top_k=None):
        probs = self.__call__(inputs, lengths)
        if top_k is not None and top_k > 0:
            top_k_preds = torch.topk(probs, top_k, dim=1)
            pred_labels = []
            for pred, prob in zip(top_k_preds.indices, top_k_preds.values):
                pred_labels.append([(self.labels[p.item()], pr.item()) for (p, pr) in zip(pred, prob)])
            return pred_labels
        if self.multilabel:
            batch_idx, label_idx = torch.where(probs > threshold)
            output = [[] for _ in range(len(inputs))]
            for batch, label in zip(batch_idx, label_idx):
                label_string = self.labels
                output[batch.item()].append(
                    (self.labels[label.item()], probs[batch, label].item())
                )
        return output