Upload folder using huggingface_hub

config.json

@@ -0,0 +1,20 @@
+{
+  "auto_map": {
+    "AutoConfig": "configuration_sykollm.SykoConfig",
+    "AutoModelForCausalLM": "modeling_sykollm.SykoSLM"
+  },
+  "bos_token_id": 2,
+  "chunk_size": 128,
+  "context_size": 1024,
+  "d_model": 768,
+  "eos_token_id": 3,
+  "intermediate_size": 3072,
+  "model_type": "sykollm",
+  "n_heads": 6,
+  "n_layers": 24,
+  "num_global_memory_tokens": 32,
+  "num_memory_tokens": 16,
+  "pad_token_id": 0,
+  "transformers_version": "5.2.0",
+  "vocab_size": 32000
+}

configuration_sykollm.py

@@ -0,0 +1,28 @@
+from transformers import PretrainedConfig
+
+class SykoConfig(PretrainedConfig):
+    model_type = "sykollm"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        d_model=768,
+        n_layers=24,
+        n_heads=6,
+        num_memory_tokens=16,
+        num_global_memory_tokens=32,
+        intermediate_size=3072,
+        chunk_size=128,
+        context_size=1024,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.num_memory_tokens = num_memory_tokens
+        self.num_global_memory_tokens = num_global_memory_tokens
+        self.intermediate_size = intermediate_size
+        self.chunk_size = chunk_size
+        self.context_size = context_size

model.safetensors

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

modeling_sykollm.py

@@ -0,0 +1,96 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from configuration_sykollm import SykoConfig
+
+
+class SykoMemoryGate(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.forget_linear = nn.Linear(d_model * 2, d_model)
+        self.update_linear = nn.Linear(d_model, d_model)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, current_context, prev_memory):
+        combined = torch.cat([current_context, prev_memory], dim=-1)
+        forget_ratio = torch.sigmoid(self.forget_linear(combined))
+        new_candidate = torch.tanh(self.update_linear(current_context))
+        new_memory = (forget_ratio * prev_memory) + ((1 - forget_ratio) * new_candidate)
+        return self.norm(new_memory)
+
+
+class SykoSmartMemoryGate(nn.Module):
+    def __init__(self, d_model, num_heads=4):
+        super().__init__()
+        self.summarizer = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
+        self.forget_linear = nn.Linear(d_model * 2, d_model)
+        self.update_linear = nn.Linear(d_model, d_model)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, full_chunk_output, global_memory_output, prev_global_memory):
+        summary, _ = self.summarizer(
+            query=global_memory_output,
+            key=full_chunk_output,
+            value=full_chunk_output
+        )
+        combined = torch.cat([summary, prev_global_memory], dim=-1)
+        forget_ratio = torch.sigmoid(self.forget_linear(combined))
+        new_candidate = torch.tanh(self.update_linear(summary))
+        new_memory = (forget_ratio * prev_global_memory) + ((1 - forget_ratio) * new_candidate)
+        return self.norm(new_memory)
+
+
+class SykoSLM(PreTrainedModel):
+    config_class = SykoConfig
+
+    def __init__(self, config: SykoConfig):
+        super().__init__(config)
+        self.d_model = config.d_model
+        self.mem_tokens = config.num_memory_tokens
+        self.g_mem_tokens = config.num_global_memory_tokens
+
+        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.pos_embedding = nn.Embedding(
+            config.context_size + config.num_memory_tokens + config.num_global_memory_tokens,
+            config.d_model
+        )
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.d_model,
+            nhead=config.n_heads,
+            dim_feedforward=config.intermediate_size,
+            batch_first=True,
+            norm_first=True
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=config.n_layers, norm=nn.LayerNorm(config.d_model))
+        self.memory_gate = SykoMemoryGate(config.d_model)
+        self.global_memory_gate = SykoSmartMemoryGate(config.d_model, num_heads=4)
+        self.fc_out = nn.Linear(config.d_model, config.vocab_size)
+
+    def forward(self, input_ids, prev_memory, global_memory, chunk_start_idx=0):
+        x = self.embedding(input_ids)
+        x_with_memory = torch.cat([global_memory, prev_memory, x], dim=1)
+        seq_len = x.size(1)
+
+        g_mem_positions = torch.arange(0, self.g_mem_tokens, device=input_ids.device)
+        mem_positions   = torch.arange(self.g_mem_tokens, self.g_mem_tokens + self.mem_tokens, device=input_ids.device)
+        word_positions  = torch.arange(
+            self.g_mem_tokens + self.mem_tokens + chunk_start_idx,
+            self.g_mem_tokens + self.mem_tokens + chunk_start_idx + seq_len,
+            device=input_ids.device
+        )
+        positions = torch.cat([g_mem_positions, mem_positions, word_positions]).unsqueeze(0)
+        x_with_memory = x_with_memory + self.pos_embedding(positions)
+
+        causal_mask = nn.Transformer.generate_square_subsequent_mask(x_with_memory.size(1), device=input_ids.device)
+        out = self.transformer(x_with_memory, mask=causal_mask)
+
+        global_memory_output = out[:, :self.g_mem_tokens, :]
+        memory_output        = out[:, self.g_mem_tokens : self.g_mem_tokens + self.mem_tokens, :]
+        token_outputs        = out[:, self.g_mem_tokens + self.mem_tokens:, :]
+
+        logits            = self.fc_out(token_outputs)
+        new_memory        = self.memory_gate(memory_output, prev_memory)
+        new_global_memory = self.global_memory_gate(token_outputs, global_memory_output, global_memory)
+
+        return logits, new_memory, new_global_memory

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}