Model save

MemoryCell.py

@@ -0,0 +1,208 @@
+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

PreTrainedRMTConfig.py

@@ -1,15 +1,21 @@
 import os
 import json
-from transformers import PretrainedConfig
+from typing import Type
+from transformers import AutoConfig, PretrainedConfig
+
+def register_to_hf_auto_config(
+    config_class: Type[PretrainedConfig],
+) -> Type[PretrainedConfig]:
+    AutoConfig.register(config_class.model_type, config_class)
+    return config_class
 
 class PreTrainedRMTConfig(PretrainedConfig):
     """
-    Recurrent Memory Transformer の設定クラス
+    Recurrent Memory Transformer configuration class
     """
     
     model_type = "rmt"
     
-    # マッピング情報を追加（設定クラスとモデルクラスの関連付け）
     auto_map = {
         "AutoModelForCausalLM": "open_r1.rmt.RecurrentMemoryTransofomer.RecurrentMemoryTransformer"
     }
@@ -45,12 +51,10 @@ class PreTrainedRMTConfig(PretrainedConfig):
             self.base_model_type = dict_config.get("model_type")
             if self.base_model_type is None:
                 raise ValueError("base_model_configにmodel_typeが指定されていません。")
-            PreTrainedRMTConfig.model_type = "rmt_" + self.base_model_type   
+            #PreTrainedRMTConfig.model_type = "rmt_" + self.base_model_type  
     """
     def __repr__(self):
         return f"PreTrainedRMTConfig(is_memory_all={self.is_memory_all}, max_n_segments={self.max_n_segments}, " \
                f"input_seg_len={self.input_seg_len}, output_seg_len={self.output_seg_len}, " \
                f"align='{self.align}', num_mem_tokens={self.num_mem_tokens})"
-    """
-
-PreTrainedRMTConfig.register_for_auto_class()
+    """

README.md

@@ -1,11 +1,9 @@
 ---
 base_model: openai-community/gpt2
-datasets: HuggingFaceFW/fineweb-edu
 library_name: transformers
 model_name: gpt2-RMT-2-mem512
 tags:
 - generated_from_trainer
-- open-r1
 - trl
 - sft
 licence: license
@@ -13,7 +11,7 @@ licence: license
 
 # Model Card for gpt2-RMT-2-mem512
 
-This model is a fine-tuned version of [openai-community/gpt2](https://huggingface.co/openai-community/gpt2) on the [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) dataset.
+This model is a fine-tuned version of [openai-community/gpt2](https://huggingface.co/openai-community/gpt2).
 It has been trained using [TRL](https://github.com/huggingface/trl).
 
 ## Quick start
@@ -29,7 +27,7 @@ print(output["generated_text"])
 
 ## Training procedure
 
-[<img src="https://raw.githubusercontent.com/wandb/assets/main/wandb-github-badge-28.svg" alt="Visualize in Weights & Biases" width="150" height="24"/>](https://wandb.ai/shin2021001-osaka-city-university/huggingface/runs/nt4l8say) 
+[<img src="https://raw.githubusercontent.com/wandb/assets/main/wandb-github-badge-28.svg" alt="Visualize in Weights & Biases" width="150" height="24"/>](https://wandb.ai/shin2021001-osaka-city-university/huggingface/runs/p1finncz) 
 
 
 This model was trained with SFT.
@@ -38,7 +36,7 @@ This model was trained with SFT.
 
 - TRL: 0.15.2
 - Transformers: 4.50.0.dev0
-- Pytorch: 2.5.1
+- Pytorch: 2.5.1+cu121
 - Datasets: 3.3.2
 - Tokenizers: 0.21.0
 

RecurrentMemoryTransformer.py

@@ -0,0 +1,171 @@
+import torch
+from transformers import PreTrainedModel, AutoModelForCausalLM, AutoConfig
+from transformers.models.auto.auto_factory import _BaseAutoModelClass
+from .MemoryCell import MemoryCell
+from .RecurrentWrapper import RecurrentWrapper
+from .PreTrainedRMTConfig import PreTrainedRMTConfig
+
+
+# @register_for_auto_class("AutoModelForCausalLM")
+class RecurrentMemoryTransformer(PreTrainedModel):
+    """
+    Recurrent Memory Transformer Model Class
+    A transformer model that processes long context in segments and retains information using memory
+    """
+    
+    config_class = PreTrainedRMTConfig
+    auto_model_class = "AutoModelForCausalLM"
+    
+    # マッピングを定義してAutoクラスが適切なモデルを見つけられるようにする
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+    
+    # AUTO_MAPを定義（モデル名からクラスへのマッピング）
+    AUTO_MAP = {
+        "AutoModelForCausalLM": "RecurrentMemoryTransformer",
+    }
+    
+    def __init__(self, config, base_model=None):
+        """
+        Initialization
+        
+        Parameters
+        ----------
+        config : PreTrainedRMTConfig
+            Model configuration
+        base_model : PreTrainedModel, optional
+            Base transformer model
+        """
+        super().__init__(config)
+        
+        # base_modelが指定されていない場合は、configから自動生成
+        if base_model is None:
+            # ベースモデルのタイプを確認
+            if not hasattr(config, "base_model_type"):
+                raise ValueError("configにbase_model_typeが指定されていません。RMTの設定にはベースモデルタイプが必要です。")
+            base_model_type = config.base_model_type
+            
+            # ベースモデル用の設定を作成
+            base_config = AutoConfig.from_pretrained(base_model_type)
+            
+            # RMT固有のパラメータを除外してベースモデルの設定を作成
+            rmt_specific_params = ['model_type', 'is_memory_all', 'max_n_segments', 'input_seg_len', 
+                                  'output_seg_len', 'align', 'num_mem_tokens', 'base_model_type']
+            for key, value in config.__dict__.items():
+                if key not in rmt_specific_params and not key.startswith('_'):
+                    setattr(base_config, key, value)
+            
+            # ベースモデルを作成
+            base_model = AutoModelForCausalLM.from_config(base_config)
+        
+        # MemoryCellとRecurrentWrapperの初期化
+        memory_cell = MemoryCell(base_model, config.num_mem_tokens)
+        self.recurrent_wrapper = RecurrentWrapper(
+            memory_cell=memory_cell,
+            is_memory_all=config.is_memory_all,
+            max_n_segments=config.max_n_segments,
+            input_seg_len=config.input_seg_len,
+            output_seg_len=config.output_seg_len,
+            align=config.align
+        )
+        
+    def get_base_model(self):
+        """
+        Get the base model
+        """
+        return self.recurrent_wrapper.memory_cell.model
+    
+    def forward(self, input_ids=None, attention_mask=None, labels=None, labels_mask=None, 
+                inputs_embeds=None, output_attentions=None, output_hidden_states=None):
+        """
+        Forward pass of the model
+        
+        Parameters
+        ----------
+        input_ids : torch.Tensor, optional
+            Input tensor
+        attention_mask : torch.Tensor, optional
+            Attention mask
+        labels : torch.Tensor, optional
+            Label tensor
+        labels_mask : torch.Tensor, optional
+            Label mask
+        inputs_embeds : torch.Tensor, optional
+            Input embeddings
+        output_attentions : bool, optional
+            Whether to output attention weights
+        output_hidden_states : bool, optional
+            Whether to output hidden states
+        """
+        forward_kwargs = {}
+        if input_ids is not None:
+            forward_kwargs["input_ids"] = input_ids
+        if labels is not None:
+            forward_kwargs["labels"] = labels
+        if attention_mask is not None:
+            forward_kwargs["attention_mask"] = attention_mask
+        if labels_mask is not None:
+            forward_kwargs["labels_mask"] = labels_mask
+        if inputs_embeds is not None:
+            forward_kwargs["inputs_embeds"] = inputs_embeds
+        if output_attentions is not None:
+            forward_kwargs["output_attentions"] = output_attentions
+        if output_hidden_states is not None:
+            forward_kwargs["output_hidden_states"] = output_hidden_states
+        
+        #forward_kwargs.update(kwargs)
+        
+        # 通常の順伝播処理
+        out = self.recurrent_wrapper.forward(**forward_kwargs)
+        """
+        # デバッグ出力を削除（または必要に応じてコメント化）
+        # print(out["loss"])
+        
+        # 分散環境で損失が二��計算されないよう、ワールドサイズで割る
+        # これは処理済みの場合は不要なので、環境変数などで制御することも可能
+        if torch.distributed.is_initialized() and "loss" in out and out["loss"] is not None:
+            # 既にDeepSpeedが処理している可能性があるため、確認が必要
+            # テスト目的で一時的に追加（実際の環境に合わせて調整が必要）
+            # world_size = torch.distributed.get_world_size()
+            # out["loss"] = out["loss"] / world_size
+            pass
+        """
+        return out
+    
+    def generate(self, **kwargs):
+        """
+        Text generation
+        """
+        return self.recurrent_wrapper.generate(**kwargs)
+    
+    def generate_with_tokenizer(self, tokenizer, input_text, **kwargs):
+        """
+        Text generation using tokenizer
+        """
+        return self.recurrent_wrapper.generate_with_tokenizer(tokenizer, input_text, **kwargs)
+    
+    def get_input_embeddings(self):
+        """
+        Get input embeddings
+        """
+        return self.get_base_model().get_input_embeddings()
+    
+    def set_input_embeddings(self, embeddings):
+        """
+        Set input embeddings
+        """
+        self.get_base_model().set_input_embeddings(embeddings)
+    
+    def get_output_embeddings(self):
+        """
+        Get output embeddings
+        """
+        return self.get_base_model().get_output_embeddings()
+    
+    def resize_token_embeddings(self, new_num_tokens):
+        """
+        Resize token embeddings
+        """
+        self.get_base_model().resize_token_embeddings(new_num_tokens)
+        return self.get_input_embeddings()
+
+RecurrentMemoryTransformer.register_for_auto_class("AutoModelForCausalLM")

RecurrentWrapper.py

@@ -0,0 +1,519 @@
+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
+

all_results.json

@@ -3,10 +3,10 @@
     "eval_samples": 100,
     "eval_samples_per_second": 376.2,
     "eval_steps_per_second": 23.56,
-    "total_flos": 5419008396361728.0,
-    "train_loss": 4.076150745858688,
-    "train_runtime": 7573.4415,
+    "total_flos": 5418484972388352.0,
+    "train_loss": 3.606253622488408,
+    "train_runtime": 424.9732,
     "train_samples": 19883,
-    "train_samples_per_second": 87.49,
-    "train_steps_per_second": 2.734
+    "train_samples_per_second": 48.742,
+    "train_steps_per_second": 1.522
 }

config.json

@@ -103,12 +103,12 @@
   "embd_pdrop": 0.1,
   "eos_token_id": 50256,
   "initializer_range": 0.02,
-  "input_seg_len": 16,
+  "input_seg_len": 512,
   "is_memory_all": false,
   "layer_norm_epsilon": 1e-05,
   "max_n_segments": 2,
   "memory_size": 512,
-  "model_type": "rmt_gpt2",
+  "model_type": "rmt",
   "n_ctx": 1024,
   "n_embd": 768,
   "n_head": 12,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:65cd08da7dc4048511a86bef339939ec8531258568d7775e32f42921c96aaab4
+oid sha256:41c0d6e17ff62620d8f534dc2766060257a8bd950e39f3902a2e65e00a21481c
 size 248915448

train_results.json

@@ -1,8 +1,8 @@
 {
-    "total_flos": 5419008396361728.0,
-    "train_loss": 4.076150745858688,
-    "train_runtime": 7573.4415,
+    "total_flos": 5418484972388352.0,
+    "train_loss": 3.606253622488408,
+    "train_runtime": 424.9732,
     "train_samples": 19883,
-    "train_samples_per_second": 87.49,
-    "train_steps_per_second": 2.734
+    "train_samples_per_second": 48.742,
+    "train_steps_per_second": 1.522
 }

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3e506da221a187e12d0f07664922d412eb372d52d46c7a3b6e4d2d2ee1a0abcd
+oid sha256:cbb45d4b8223f141e7950f15066fdb3796697a543d5274ffce9e5110eceddf62
 size 7352