Upload 2 files

configuration_sykoomni.py

@@ -0,0 +1,43 @@
+from transformers import PretrainedConfig
+
+class SykoConfig(PretrainedConfig):
+    model_type = "sykoomni"
+
+    def __init__(
+        self,
+        text_vocab_size=32000,
+        image_vocab_size=8192,
+        audio_vocab_size=1024,
+        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,
+        causal_reasoning_layers=2,
+        vision_hidden_size=768,
+        audio_hidden_size=384,
+        max_image_tokens=256,
+        max_audio_tokens=500,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.text_vocab_size          = text_vocab_size
+        self.image_vocab_size         = image_vocab_size
+        self.audio_vocab_size         = audio_vocab_size
+        self.vocab_size               = text_vocab_size + image_vocab_size + audio_vocab_size + 10
+        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
+        self.causal_reasoning_layers  = causal_reasoning_layers
+        self.vision_hidden_size       = vision_hidden_size
+        self.audio_hidden_size        = audio_hidden_size
+        self.max_image_tokens         = max_image_tokens
+        self.max_audio_tokens         = max_audio_tokens

modeling_sykoomni.py

@@ -0,0 +1,192 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, SiglipVisionModel, WhisperModel
+from .configuration_sykoomni 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))
+        return self.norm((forget_ratio * prev_memory) + ((1 - forget_ratio) * new_candidate))
+
+
+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 * 3, 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, context=None):
+        summary, _ = self.summarizer(query=global_memory_output,
+                                     key=full_chunk_output, value=full_chunk_output)
+        ctx = (context.mean(dim=1, keepdim=True).expand_as(summary)
+               if context is not None else torch.zeros_like(summary))
+        combined      = torch.cat([summary, ctx, prev_global_memory], dim=-1)
+        forget_ratio  = torch.sigmoid(self.forget_linear(combined))
+        new_candidate = torch.tanh(self.update_linear(summary))
+        return self.norm((forget_ratio * prev_global_memory) + ((1 - forget_ratio) * new_candidate))
+
+
+class SykoCausalHead(nn.Module):
+    def __init__(self, d_model, vocab_size, num_layers=2):
+        super().__init__()
+        layers = []
+        for _ in range(num_layers):
+            layers.extend([nn.Linear(d_model, d_model), nn.GELU(), nn.LayerNorm(d_model)])
+        layers.append(nn.Linear(d_model, vocab_size))
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class SykoModalEmbedding(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.modal_embed = nn.Embedding(4, d_model)
+
+    def forward(self, x, modal_type):
+        return x + self.modal_embed(torch.tensor(modal_type, device=x.device))
+
+
+class SykoOmni(PreTrainedModel):
+    config_class       = SykoConfig
+    _tied_weights_keys = []
+
+    def __init__(self, config: SykoConfig):
+        super().__init__(config)
+        if not hasattr(self, "all_tied_weights_keys"):
+            self.all_tied_weights_keys = {}
+        self.d_model      = config.d_model
+        self.mem_tokens   = config.num_memory_tokens
+        self.g_mem_tokens = config.num_global_memory_tokens
+
+        self.text_embedding  = nn.Embedding(config.text_vocab_size, config.d_model)
+        self.image_embedding = nn.Embedding(config.image_vocab_size, config.d_model)
+        self.audio_embedding = nn.Embedding(config.audio_vocab_size, config.d_model)
+
+        self.pos_embedding = nn.Embedding(
+            config.context_size + config.num_memory_tokens +
+            config.num_global_memory_tokens + 50, 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.text_head  = nn.Linear(config.d_model, config.text_vocab_size)
+        self.image_head = nn.Linear(config.d_model, config.image_vocab_size)
+        self.audio_head = nn.Linear(config.d_model, config.audio_vocab_size)
+
+        self.causal_head     = SykoCausalHead(config.d_model, config.text_vocab_size,
+                                               config.causal_reasoning_layers)
+        self.modal_embedding = SykoModalEmbedding(config.d_model)
+
+        self.vision_encoder = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224")
+        for p in self.vision_encoder.parameters():
+            p.requires_grad = False
+        self.vision_proj = nn.Sequential(
+            nn.Linear(config.vision_hidden_size, config.d_model),
+            nn.GELU(), nn.LayerNorm(config.d_model))
+
+        self.audio_encoder = WhisperModel.from_pretrained("openai/whisper-tiny").encoder
+        for p in self.audio_encoder.parameters():
+            p.requires_grad = False
+        self.audio_proj = nn.Sequential(
+            nn.Linear(config.audio_hidden_size, config.d_model),
+            nn.GELU(), nn.LayerNorm(config.d_model))
+
+    def _adjust_tied_keys_with_tied_pointers(self, missing_keys):
+        pass
+
+    def _token_ids_to_embeddings(self, input_ids, text_vocab_size, image_vocab_size):
+        batch, seq = input_ids.shape
+        embeddings = torch.zeros(batch, seq, self.d_model, device=input_ids.device)
+
+        text_mask = input_ids < text_vocab_size
+        if text_mask.any():
+            safe = input_ids.clone()
+            safe[~text_mask] = 0
+            embeddings[text_mask] = self.text_embedding(safe)[text_mask]
+
+        img_start = text_vocab_size
+        img_end   = text_vocab_size + image_vocab_size
+        img_mask  = (input_ids >= img_start) & (input_ids < img_end)
+        if img_mask.any():
+            ids = input_ids.clone()
+            ids[~img_mask] = 0
+            embeddings[img_mask] = self.image_embedding(
+                (ids - img_start).clamp(0, image_vocab_size - 1))[img_mask]
+
+        aud_mask = input_ids >= img_end
+        if aud_mask.any():
+            ids = input_ids.clone()
+            ids[~aud_mask] = 0
+            embeddings[aud_mask] = self.audio_embedding(
+                (ids - img_end).clamp(0, self.config.audio_vocab_size - 1))[aud_mask]
+
+        return embeddings
+
+    def forward(self, input_ids, prev_memory, global_memory,
+                chunk_start_idx=0, pixel_values=None, audio_features=None,
+                generation_mode=None):
+        cfg = self.config
+        x   = self._token_ids_to_embeddings(input_ids, cfg.text_vocab_size, cfg.image_vocab_size)
+        x   = self.modal_embedding(x, modal_type=0)
+
+        if pixel_values is not None:
+            with torch.no_grad():
+                vis_feats = self.vision_encoder(pixel_values=pixel_values).last_hidden_state
+            x = torch.cat([self.modal_embedding(self.vision_proj(vis_feats), 1), x], dim=1)
+
+        if audio_features is not None:
+            with torch.no_grad():
+                aud_feats = self.audio_encoder(audio_features).last_hidden_state
+            x = torch.cat([self.modal_embedding(self.audio_proj(aud_feats), 2), x], dim=1)
+
+        seq_len       = x.size(1)
+        x_with_memory = torch.cat([global_memory, prev_memory, x], dim=1)
+
+        g_pos = torch.arange(0, self.g_mem_tokens, device=input_ids.device)
+        m_pos = torch.arange(self.g_mem_tokens, self.g_mem_tokens + self.mem_tokens, device=input_ids.device)
+        t_pos = 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)
+        pos_ids = torch.cat([g_pos, m_pos, t_pos]).clamp(0, self.pos_embedding.num_embeddings - 1)
+        x_with_memory = x_with_memory + self.pos_embedding(pos_ids.unsqueeze(0))
+
+        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)
+
+        gmo = out[:, :self.g_mem_tokens, :]
+        mo  = out[:, self.g_mem_tokens:self.g_mem_tokens + self.mem_tokens, :]
+        to_ = out[:, self.g_mem_tokens + self.mem_tokens:, :]
+
+        if generation_mode == 'image':
+            logits = self.image_head(to_)
+        elif generation_mode == 'audio':
+            logits = self.audio_head(to_)
+        else:
+            logits = self.text_head(to_)
+
+        non_text_len = seq_len - input_ids.size(1)
+        return (logits, logits,
+                self.memory_gate(mo, prev_memory),
+                self.global_memory_gate(to_, gmo, global_memory, context=to_),
+                non_text_len)