MemoryCell.py

import math
import torch
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
from transformers import PreTrainedModel

from .PreTrainedRMTConfig import PreTrainedRMTConfig

class MemoryCell(torch.nn.Module):
    """Holds memory tensors.
    Replicates memory tensor for each batch size.
    Adds memory tokens to the input tensor and returns that tensor.
    Processes the model output and returns a new memory state.

    Parameters
    ----------
    torch : _type_
        _description_
    """

    def __init__(self, base_model, num_mem_tokens):
        super().__init__()
        self.model = base_model
        self.create_memory(num_mem_tokens)
        self.config = base_model.config
        
        # token_type_embeddingsの追加
        #self.token_type_embeddings = torch.nn.Embedding(2, getattr(self.model.config, "n_embd", self.model.config.hidden_size))

    def create_memory(self, num_mem_tokens):
        """Randomly initializes an embedding matrix (tensor) for memory tokens and registers it for gradient computation.
           Sets read and write positions for memory tokens.

        Parameters
        ----------
        num_mem_tokens : _type_
            Number of memory tokens.
        """
        self.read_memory_position = range(num_mem_tokens)
        self.write_memory_position = range(-num_mem_tokens, 0)
        
        self.num_mem_tokens = num_mem_tokens
        embeddings = self.model.get_input_embeddings()
        memory_dim = getattr(self.model.config, "n_embd", self.model.config.hidden_size)
        memory_weights = (
            torch.randn((num_mem_tokens, memory_dim))# * embeddings.weight.data.std()
        )
        
        self.register_parameter(
            "memory", torch.nn.Parameter(memory_weights, requires_grad=True)
        )

    def set_memory(self, input_shape):
        """Replicates memory tensor for each batch size

        Parameters
        ----------
        input_shape : _type_
            _description_

        Returns
        -------
        _type_
            Replicated memory tensor. (batch_size, num_mem_tokens, memory_dim)
        """
        memory = self.memory.repeat(
            input_shape[0], 1, 1
        )  # 　メモリテンソルをバッチサイズ分だけ複製する
        return memory  # (batch_size, num_mem_tokens, memory_dim)

    def forward(self, input_ids, memory_state=None, **kwargs):
        """Performs inference.

        Parameters
        ----------
        input_ids : torch.Tensor
            Input tensor.
        memory_state : torch.Tensor, optional
            Memory tensor, by default None (num_mem_tokens, memory_dim)

        Returns
        -------
        tuple(tuple, torch.Tensor)
            out : tuple
                Model output.
            new_memory_state : torch.Tensor
                New memory state.
        """
        if memory_state is None:
            # メモリテンソルをバッチサイズ分だけ複製する
            memory_state = self.set_memory(input_ids.shape)

        # メモリトークンを入力テンソルに追加し、そのテンソルを返す
        seg_kwargs = self.process_input(input_ids, memory_state, **kwargs)
        out = self.model(**seg_kwargs)
        #print(out)

        # モデルの出力を処理し、新しいメモリ状態を返す
        out, new_memory_state = self.process_output(out, **kwargs)

        return out, new_memory_state

    def process_input(self, input_ids, memory_state, **kwargs):
        """Adds memory tokens to the input tensor and returns that tensor

        Parameters
        ----------
        input_ids : _type_
            Input tensor.
        memory_state : _type_
            Memory tensor.

        Returns
        -------
        _type_
            Input tensor with added memory tokens. (batch_size, seq_len, hidden_size)
        """
        seg_kwargs = dict(**kwargs)

        inputs_embeds = kwargs.get("inputs_embeds")
        if inputs_embeds is None:
            inputs_embeds = self.model.get_input_embeddings()(input_ids)
        if inputs_embeds.shape[0] != memory_state.shape[0]: # バッチサイズが異なる場合
            memory_state = self.set_memory(inputs_embeds.shape)
        
        # メモリトークンを入力テンソルに追加
        inputs_embeds = torch.cat(
            [memory_state, inputs_embeds, memory_state], dim=1
        ).to(input_ids.device)
        """
        # token_type_idsの生成
        token_type_ids = torch.zeros_like(inputs_embeds[:, :, 0], dtype=torch.long)
        token_type_ids[:, self.num_mem_tokens:-self.num_mem_tokens] = 1

        # token_type_embeddingsの追加と入力の更新
        token_type_embeds = self.token_type_embeddings(token_type_ids)
        inputs_embeds = inputs_embeds + token_type_embeds
        """

        seg_kwargs["input_ids"] = None
        seg_kwargs["inputs_embeds"] = inputs_embeds
        if kwargs.get("attention_mask") is not None:
            seg_kwargs["attention_mask"] = self.pad_attention_mask(
                kwargs["attention_mask"], inputs_embeds.shape
            )
        seg_kwargs["output_hidden_states"] = True
        
        # Positional Embeddings
        pos_mem1 = torch.arange(self.num_mem_tokens, device=input_ids.device)
        pos_mem2 = torch.arange(self.num_mem_tokens, self.num_mem_tokens * 2, device=input_ids.device)
        pos_seg = torch.arange(self.num_mem_tokens * 2, self.num_mem_tokens * 2 + input_ids.shape[1], device=input_ids.device)
        pos = torch.cat([pos_mem1, pos_seg, pos_mem2], dim=0)
        pos = pos.unsqueeze(0).expand(input_ids.shape[0], -1)
        seg_kwargs["position_ids"] = pos
        
        return seg_kwargs

    def pad_attention_mask(self, attention_mask, shape):
        if self.num_mem_tokens in {0, None}:
            return attention_mask
        else:
            attention_mask = torch.cat(
                [
                    torch.ones(
                        shape[0], self.num_mem_tokens, device=attention_mask.device
                    ),
                    attention_mask,
                    torch.ones(
                        shape[0], self.num_mem_tokens, device=attention_mask.device
                    ),
                ],
                dim=1,
            )
            return attention_mask

    def compute_logpi(mean, stddev, action):
        a1 =-0.5 * torch.log(2*torch.fill(stddev.shape, math.pi))
        a2 = -torch.log(stddev)
        a3 = -0.5 * (((action - mean) / stddev) ** 2)
        return a1 + a2 + a3

    def process_output(self, model_outputs, **kwargs):
        if self.num_mem_tokens not in {0, None}:
            out = CausalLMOutputWithCrossAttentions()
            memory_state = model_outputs.hidden_states[-1][:, -self.num_mem_tokens :]
            out["logits"] = model_outputs.logits[
                :, self.num_mem_tokens : -self.num_mem_tokens
            ]

            if kwargs.get("output_hidden_states"):
                out["hidden_states"] = [
                    lh[:, self.num_mem_tokens : -self.num_mem_tokens]
                    for lh in model_outputs.hidden_states
                ]
            if kwargs.get("output_attentions"):
                out["attentions"] = model_outputs["attentions"]
        else:
            memory_state = None
            out = model_outputs

        return out, memory_state

    def generate(self, input_ids, memory_state, attention_mask, **generate_kwargs):
        if memory_state is None:
            memory_state = self.set_memory(input_ids.shape)

        seg_kwargs = self.process_input(input_ids, memory_state, attention_mask=attention_mask)
        out = self.model.generate(inputs_embeds=seg_kwargs['inputs_embeds'], attention_mask=seg_kwargs['attention_mask'], **generate_kwargs)
        return out