RecurrentWrapper.py

import math
import torch
from torch.nn import CrossEntropyLoss
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
from .PreTrainedRMTConfig import PreTrainedRMTConfig
from .MemoryCell import MemoryCell
from torch.nn.utils.rnn import pad_sequence
from transformers import PreTrainedModel

class RecurrentWrapper(torch.nn.Module):
    #config_class = PreTrainedRMTConfig
    
    def __init__(
        self, 
        memory_cell: MemoryCell, 
        is_memory_all: bool, 
        max_n_segments: int, 
        input_seg_len: int, 
        output_seg_len: int, 
        align: str = "left"):
        
        super().__init__()
        self.memory_cell:MemoryCell = memory_cell
        self.is_memory_all = is_memory_all # Whether to share memory state between segments
        self.memory_state: torch.Tensor = None # Memory state
        self.config = memory_cell.config # Model configuration
        self.max_n_segments = max_n_segments # Maximum number of segments for backpropagation
        self.input_seg_len = input_seg_len # Segment size
        self.output_seg_len = output_seg_len
        self.align = align # Segment alignment default: left

    def forward(
        self,
        input_ids,
        labels=None,
        labels_mask=None,
        inputs_embeds=None,
        attention_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        **kwargs
    ):
        """Performs inference.

        Parameters
        ----------
        input_ids : torch.Tensor
            Input tensor. (batch_size, seq_len * n_segments)
        labels : _type_, torch.Tensor
           Input tensor. (batch_size, seq_len * n_segments)

        Returns
        ----------
        dict
            "loss" : torch.Tensor
                Loss value.
            "logits" : torch.Tensor
                Model output.
            "out[f"{key}_{seg_num}"]" : torch.Tensor
                Output for each segment.
        """
        if self.memory_state is not None:
            if self.is_memory_all is False:
                self.memory_state = None
            else :
                self.memory_state.detach()  # メモリ状態の勾配を計算しないようにする

        # 入力テンソルをセグメント単位に分割する。 (セグメントは1ステップでモデルに渡される入力のサブセット)
        segmented = self.segment(
            self.input_seg_len,
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
        )

        cell_outputs = []  # 各セグメントの出力を保存するリスト
        for seg_num, segment in enumerate(segmented):
            cell_out, self.memory_state = self.memory_cell(
                **segment, memory_state=self.memory_state, **kwargs
            )
            cell_outputs.append(cell_out)
            a = self.manage_gradients(
                self.memory_state, seg_num, len(segmented)
            )  # メモリ状態の勾配計算を制御する
            #print(seg_num, a)

        out = self.process_outputs(
            cell_outputs,
            labels=labels,
            labels_mask=labels_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )
        return out

    def log(self, t, eps = 1e-20):
        return torch.log(t.clamp(min = eps))

    def gumbel_noise(self, t):
        noise = torch.zeros_like(t).uniform_(0, 1)
        return -self.log(-self.log(noise))

    def gumbel_sample(self, t, temperature = 1., dim = -1):
        return ((t / max(float(temperature), float(1e-10))) + self.gumbel_noise(t)).argmax(dim = dim)

    def top_k(self, logits, thres = 0.9):
        k = math.ceil((1 - thres) * logits.shape[-1])
        val, ind = torch.topk(logits, k)
        probs = torch.full_like(logits, float('-inf'))
        probs.scatter_(1, ind, val)
        return probs

    def segment(self, seg_len, **kwargs):
        """
        Segments input tensors and adjusts their size. Returns a list of dicts.

        Parameters
        ----------
        **kwargs : dict
            Tensors to be segmented.
            Specify tensors that need to be split in keyword argument format.
            Example: segment(input_ids=tensor1, attention_mask=tensor2)

        Returns
        -------
        segments : list of dict
            List of dictionaries containing segmented tensors.
            Example: [{'input_ids': segment1, 'attention_mask': segment1}, {'input_ids': segment2, 'attention_mask': segment2}, ...]

        Notes
        -----
        - This function uses the `self.split_tensor` method, so `self` must implement it.
        - Each tensor is split in a specific way by `self.split_tensor`. The same keys are stored with the same order of indices.
        """
        segments = []  # 各セグメントを保存するリストを初期化
        for k, tensor in kwargs.items():  # keyで繰り返し
            if tensor is not None:
                k_segments = self.split_tensor(
                    tensor, seg_len
                )  # 2次元テンソルを分割し、セグメント化
                for s, k_seg in enumerate(k_segments):
                    if s < len(segments):
                        segments[s][k] = k_seg
                    else:
                        segments.append({k: k_seg}) # 新たな辞書 {k: k_seg} を作成し、segments リストに追加します。

        return segments

    def split_tensor(self, tensor, seg_len):
        if self.align in {"left", None}:
            split_inds = list(range(0, tensor.shape[1], seg_len)) + [
                tensor.shape[1]
            ]
            segments = [
                tensor[:, start:end] for (start, end) in zip(split_inds, split_inds[1:])
            ]
        elif self.align in {"right", None}:
            split_inds = (list(range(tensor.shape[1], 0, -seg_len)) + [0])[::-1]
            segments = [
                tensor[:, start:end] for (start, end) in zip(split_inds, split_inds[1:])
            ]
        elif self.align == "center":
            n_seg = math.ceil(tensor.shape[1] / seg_len)
            segments = torch.chunk(tensor, n_seg, dim=1)
        else:
            split_inds = list(range(0, tensor.shape[1], seg_len)) + [
                tensor.shape[1]
            ]
            segments = [
                tensor[:, start:end] for (start, end) in zip(split_inds, split_inds[1:])
            ]
        return segments

    def process_outputs(self, cell_outputs, **kwargs):
        """Calculates loss for a list of outputs. Also concatenates and returns logits.

        Parameters
        ----------
        cell_outputs : list of torch.Tensor
            List containing outputs from each segment.

        Returns
        -------
        dict
            "loss" : torch.Tensor
                Loss value.
            "logits" : torch.Tensor
                Model output.
            "out[f"{key}_{seg_num}"]" : torch.Tensor
                Output for each segment.
        """
        out = CausalLMOutputWithCrossAttentions()
        full_logits = torch.cat(
            [o.logits for o in cell_outputs], dim=1
        )  # セグメントごとのlogitsを結合する (batch_size, seq_len * seg_len, vocab_size)

        if kwargs.get("output_hidden_states"):
            full_hidden_states = tuple(
                [
                    torch.cat(layer_hs, dim=1)
                    for layer_hs in zip(*[o.hidden_states for o in cell_outputs])
                ]
            )

        labels = kwargs.get("labels")
        if labels is not None:  # ラベルがある場合のみlossを計算する
            
            shift_labels = labels[..., 1:].contiguous() # DataSetでシフトされない場合
            shift_logits = full_logits[..., :-1, :].contiguous()#  DataSetでシフトされない場合
            #shift_labels = labels.contiguous() # DataSetでシフトされる場合
            #shift_logits = full_logits.contiguous() # DataSetでシフトされる場合
            
            flat_labels = shift_labels.view(
                -1
            )  # バッチとセグメントの次元を結合して1次元にする (batch_size * (seq_len-1) * seg_len)
            flat_logits = shift_logits.view(
                -1, shift_logits.size(-1)
            )  # バッチとセグメントの次元を結合して1次元にする (batch_size * (seq_len-1) * seg_len, vocab_size)

            loss_fct = CrossEntropyLoss()
            labels_mask = kwargs.get("labels_mask")
            if labels_mask is not None:
                shift_mask = labels_mask[..., :-1].contiguous()

                flat_labels = flat_labels[shift_mask.view(-1)]
                flat_logits = flat_logits[shift_mask.view(-1)]
            out["loss"] = loss_fct(flat_logits, flat_labels)
        else:
            out["loss"] = 0
            print("labels is None")

        out["logits"] = full_logits
        segment_keys = ["loss", "logits"]
        if kwargs.get("output_attentions"):
            segment_keys.append("attentions")
        if kwargs.get("output_hidden_states"):
            segment_keys.append("hidden_states")
            out["hidden_states"] = full_hidden_states

        for seg_num, o in enumerate(cell_outputs):
            for key, value in o.items():
                if any([sk in key for sk in segment_keys]):
                    out[f"{key}_{seg_num}"] = value

        return out

    def manage_gradients(self, memory_state, seg_num, seg_len):
        """Controls gradient calculation for memory state

        Parameters
        ----------
        memory_state : torch.Tensor
            Memory state. (batch_size, num_mem_tokens, memory_dim)
        seg_num : int
            Number of the segment currently being processed.

        Returns
        ----------
        bool
            Whether to calculate gradients. True: calculate gradients, False: do not calculate gradients
        """

        # max_n_segments: 処理できる最大セグメント数を示すパラメータです。この値を使って、必要に応じてメモリの更新を決定します。

        # seg_numが0の時はReccurentでない時なので勾配は計算する。
        # 最後のほうのセグメントは勾配を計算する。
        if seg_num == 0 or self.max_n_segments in {-1, None} or seg_len - seg_num <= self.max_n_segments:
            self.memory_state = memory_state  # Retain gradients
            return True
        else:
            self.memory_state = memory_state.detach()  # Detach to stop gradient tracking
            return False

    def generate_groq(
        self,
        input_ids,
        max_length=25,
        temperature=1.0,
        top_k=None,
        top_p=None,
        do_sample=True,
        pad_token_id=None,
        eos_token_id=None,
        **kwargs
    ):
        """
        Generate new tokens based on the input sequence.

        Parameters
        ----------
        input_ids : torch.Tensor
            Initial input sequence. Shape: (batch_size, seq_len)
        max_length : int
            Maximum number of tokens to generate (including initial sequence length).
        temperature : float, default 1.0
            Temperature parameter for sampling. Lower values make it more deterministic.
        top_k : int, optional
            Used to sample from top k tokens.
        top_p : float, optional
            Used to filter tokens based on cumulative probability p.
        do_sample : bool, default True
            If True, use probabilistic sampling. If False, use greedy decoding.
        pad_token_id : int, optional
            ID of the padding token.
        eos_token_id : int, optional
            ID of the end-of-sequence token.
        **kwargs : dict
            Additional arguments passed to MemoryCell.

        Returns
        -------
        torch.Tensor
            Generated token sequence. Shape: (batch_size, generated_seq_len)
        """
        # 初期の入力シーケンスを処理
        segmented = self.segment(self.input_seg_len, input_ids=input_ids)
        memory_state = None
        for segment in segmented:
            cell_out, memory_state = self.memory_cell(
                **segment, memory_state=memory_state, **kwargs
            )

        # 生成ループ
        output_ids = input_ids
        while output_ids.shape[1] < max_length:
            # 最後のトークンを input_ids として使用
            last_token = output_ids[:, -1:]
            # MemoryCell に渡す
            cell_out, memory_state = self.memory_cell(
                input_ids=last_token, memory_state=memory_state, **kwargs
            )
            # logits を取得（最後のトークンの logits）
            logits = cell_out.logits[:, -1, :]
            # 次のトークンをサンプリング
            next_token = self.sample_next_token(
                logits, temperature, top_k, top_p, do_sample
            )
            # 出力シーケンスに追加
            output_ids = torch.cat([output_ids, next_token], dim=1)
            # 終了条件をチェック
            if eos_token_id is not None and next_token.item() == eos_token_id:
                break

        return output_ids

    def sample_next_token(self, logits, temperature=1, top_k=50, top_p=0.9, do_sample=False):
        """
        logits から次のトークンをサンプリングする。

        Parameters
        ----------
        logits : torch.Tensor
            トークンの予測スコア。形状: (batch_size, vocab_size)
        temperature : float
            サンプリング時の温度パラメータ。
        top_k : int, optional
            上位 k トークンからサンプリングする場合に使用。
        top_p : float, optional
            累積確率 p に基づいてトークンをフィルタリングする場合に使用。
        do_sample : bool
            True の場合、確率的サンプリングを使用。False の場合、貪欲法を使用。

        Returns
        -------
        torch.Tensor
            サンプリングされたトークン。形状: (batch_size, 1)
        """
        if do_sample:
            if temperature != 1.0:
                logits = logits / temperature
            if top_k is not None:
                logits = self.top_k_groq(logits, top_k)
            if top_p is not None:
                logits = self.top_p(logits, top_p)
            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
        else:
            next_token = torch.argmax(logits, dim=-1, keepdim=True)
        return next_token

    def top_k_groq(self, logits, k):
        """
        上位 k トークンのみを考慮するように logits をフィルタリングする。

        Parameters
        ----------
        logits : torch.Tensor
            トークンの予測スコア。形状: (batch_size, vocab_size)
        k : int
            上位 k トークンを選択。

        Returns
        -------
        torch.Tensor
            フィルタリングされた logits。形状: (batch_size, vocab_size)
        """
        values, indices = torch.topk(logits, k, dim=-1)
        min_values = values[:, -1].unsqueeze(-1).expand_as(logits)
        return torch.where(
            logits >= min_values, logits, torch.full_like(logits, float('-inf'))
        )

    def top_p(self, logits, p):
        """
        累積確率 p に基づいてトークンをフィルタリングする。

        Parameters
        ----------
        logits : torch.Tensor
            トークンの予測スコア。形状: (batch_size, vocab_size)
        p : float
            累積確率の閾値。

        Returns
        -------
        torch.Tensor
            フィルタリングされた logits。形状: (batch_size, vocab_size)
        """
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_indices_to_remove = cumulative_probs > p
        sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
        sorted_indices_to_remove[:, 0] = 0
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits.scatter_(1, indices_to_remove, float('-inf'))
        return logits
    
    def generate_default(self, input_ids, attention_mask = None, **generate_kwargs):
        memory_state = None
        segmented = self.segment(self.input_seg_len, input_ids=input_ids, attention_mask=attention_mask)

        for seg_num, segment in enumerate(segmented[:-1]):
            cell_out, memory_state = self.memory_cell(**segment, memory_state=memory_state)

        final_segment = segmented[-1]
        out = self.memory_cell.generate(**final_segment, memory_state=memory_state, **generate_kwargs)

        return out

    def generate(self, input_ids:torch.Tensor, **generate_kwargs):
        with torch.no_grad():
            if self.is_memory_all is False:
                self.memory_state = None
            elif self.memory_state is not None:
                self.memory_state.detach()  # メモリ状態の勾配を計算しないようにする

            # 入力テンソルをセグメント化してサイズを調整　return: [{'input_ids': 分割1, 'attention_mask': 分割1}, {'input_ids': 分割2, 'attention_mask': 分割2}, ...]
            segmented = self.segment(self.input_seg_len, input_ids=input_ids)

            for seg_num, segment in enumerate(segmented[:-1]):  # 最後のセグメント以外
                # メモリセルに入力テンソルを渡し、出力と新しいメモリ状態を取得
                cell_out, self.memory_state = self.memory_cell(
                    **segment, memory_state=self.memory_state, output_hidden_states=True
                )
                
            curr_segment = segmented[-1]
            """
            outs = []
            for i in range(math.ceil(generate_kwargs["max_length"] / self.input_seg_len)):
                out = self.memory_cell.generate(
                    **curr_segment, 
                    memory_state=self.memory_state, 
                    max_length=min(generate_kwargs["max_length"] - i * self.input_seg_len, self.input_seg_len - curr_segment["input_ids"].shape[-1]), 
                    **generate_kwargs)
                outs.append(out)
                
            for out in outs:
                for key, value in out.items():
                    curr_segment[key] = torch.cat((curr_segment[key], value), dim = -1)
                self.memory_state = out["memory_state"]
            """

            output_ids = None
            if generate_kwargs.get("max_length") is None:
                length = generate_kwargs.get("max_new_tokens", 25)
            else:
                length = generate_kwargs.get("max_length") - curr_segment["input_ids"].shape[-1]

            for ind in range(length):
                # メモリセルに入力テンソルを渡し、出力と新しいメモリ状態を取得
                out, next_memories = self.memory_cell(**curr_segment, memory_state=self.memory_state, output_hidden_states=True)
                logits = out["logits"][:,-1] # (batch_size, vocab_size)
                sampled = self.sample_next_token(logits, temperature = generate_kwargs.get("temperature", 1), top_k = generate_kwargs.get("top_k", 0.9), top_p = generate_kwargs.get("top_p", 0.9), do_sample = generate_kwargs.get("do_sample", False)) # サンプリング (batch_size, 1)
                #filtered_logits = self.top_k(logits, generate_kwargs.get("top_k", 0.9)) # トップkの確率を取得
                #sampled = self.gumbel_sample(filtered_logits, temperature = generate_kwargs.get("temperture", 1)).unsqueeze(1) # サンプリング (batch_size, 1)
                
                output_ids = sampled if output_ids is None else torch.cat((output_ids, sampled), dim = 1)
                
                curr_segment["input_ids"] = torch.cat((curr_segment["input_ids"], sampled), dim = -1) # セグメントにサンプリングされたトークンを追加 (batch_size, seq_len)
                #curr_segment["attention_mask"] = torch.cat((curr_segment["attention_mask"], torch.ones_like(sampled)), dim = -1) # セグメントのアテンションマスクを更新

                if curr_segment["input_ids"].shape[-1] > self.input_seg_len: # セグメントサイズを超えた場合
                    for key, value in curr_segment.items():
                        curr_segment[key] = value[:, -1:] # セグメントサイズに切り詰める
                    self.memory_state = next_memories # メモリ状態を更新

            return output_ids

    def generate_with_tokenizer(self, tokenizer, input_text, **generate_kwargs):
        if isinstance(input_text, str):
            tok = tokenizer(input_text, return_tensors="pt")
            tok["input_ids"] = tok["input_ids"]
            tok["attention_mask"] = tok["attention_mask"]
        else:
            tok = tokenizer(input_text)
            for k, v in tok.items():
                pd = tokenizer.pad_token_id if k != 'attention_mask' else 0
                tok[k] = pad_sequence([torch.tensor(o) for o in v], padding_value=pd, padding_side="left").T
                
        output_ids = self.generate(tok["input_ids"], **generate_kwargs)
        
        if isinstance(input_text, str):
            return tokenizer.decode(torch.cat((tok["input_ids"][0], output_ids[0]), dim=0), skip_special_tokens=True)
        else:
            return tokenizer.batch_decode(torch.cat((tok["input_ids"], output_ids), dim=-1), skip_special_tokens=True)

    def can_generate(self):
        return True