Create modeling_gclm.py

modeling_gclm.py

@@ -0,0 +1,195 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+
+
+# ============================================================
+# Configuration class is assumed to live in configuration_gclm.py
+# ============================================================
+
+# Expected fields in GCLMConfig:
+# - vocab_size
+# - d_model
+# - n_layers
+# - max_seq_len
+# - local_kernel_size
+# - global_kernel_size
+# - fft_size
+# - use_global_every_n_layers
+# - layer_norm_eps
+
+
+# ============================================================
+# Global FFT Convolution
+# ============================================================
+
+class GlobalConv1D(nn.Module):
+    def __init__(self, d_model, kernel_size, fft_size):
+        super().__init__()
+        self.kernel = nn.Parameter(torch.randn(d_model, kernel_size) * 0.01)
+        self.kernel_size = kernel_size
+        self.fft_size = fft_size
+
+    def forward(self, x):
+        # x: [B, C, T]
+        B, C, T = x.shape
+        K = min(self.kernel_size, T)
+
+        overlap = K - 1
+        block = self.fft_size - overlap
+
+        x = F.pad(x, (overlap, 0))
+        k = self.kernel[:, :K]
+        k = F.pad(k, (0, self.fft_size - K))
+
+        k_f = torch.fft.rfft(k, n=self.fft_size)
+
+        outs = []
+        pos = 0
+        while pos < T:
+            seg = x[..., pos:pos + self.fft_size]
+            if seg.shape[-1] < self.fft_size:
+                seg = F.pad(seg, (0, self.fft_size - seg.shape[-1]))
+
+            y = torch.fft.irfft(
+                torch.fft.rfft(seg, n=self.fft_size) * k_f.unsqueeze(0),
+                n=self.fft_size
+            )
+            outs.append(y[..., overlap:overlap + block])
+            pos += block
+
+        return torch.cat(outs, dim=-1)[..., :T]
+
+
+# ============================================================
+# Local Convolution
+# ============================================================
+
+class LocalConv1D(nn.Module):
+    def __init__(self, d_model, k):
+        super().__init__()
+        self.k = k
+        self.dw = nn.Conv1d(d_model, d_model, k, groups=d_model)
+        self.pw = nn.Conv1d(d_model, d_model, 1)
+
+    def forward(self, x):
+        x = F.pad(x, (self.k - 1, 0))
+        return self.pw(F.relu(self.dw(x)))
+
+
+# ============================================================
+# GCLM Block
+# ============================================================
+
+class GCLMBlock(nn.Module):
+    def __init__(self, config, use_global):
+        super().__init__()
+        self.use_global = use_global
+
+        self.ln1 = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+        self.local = LocalConv1D(
+            config.d_model,
+            config.local_kernel_size
+        )
+
+        if use_global:
+            self.ln2 = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+            self.global_conv = GlobalConv1D(
+                config.d_model,
+                config.global_kernel_size,
+                config.fft_size
+            )
+
+        self.ln3 = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+        self.ff = nn.Sequential(
+            nn.Linear(config.d_model, config.d_model * 4),
+            nn.GELU(),
+            nn.Linear(config.d_model * 4, config.d_model),
+        )
+
+    def forward(self, x):
+        x = x + self.local(self.ln1(x).transpose(1, 2)).transpose(1, 2)
+        if self.use_global:
+            x = x + self.global_conv(self.ln2(x).transpose(1, 2)).transpose(1, 2)
+        return x + self.ff(self.ln3(x))
+
+
+# ============================================================
+# Base GCLM Model
+# ============================================================
+
+class GCLMModel(PreTrainedModel):
+    config_class = None  # set by AutoConfig
+    base_model_prefix = "gclm"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.emb = nn.Embedding(config.vocab_size, config.d_model)
+        self.pos = nn.Embedding(config.max_seq_len, config.d_model)
+
+        self.layers = nn.ModuleList([
+            GCLMBlock(
+                config,
+                use_global=(i % config.use_global_every_n_layers == 0)
+            )
+            for i in range(config.n_layers)
+        ])
+
+        self.ln = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+
+        self.post_init()
+
+    def forward(self, input_ids):
+        B, T = input_ids.shape
+        pos = torch.arange(T, device=input_ids.device)
+
+        h = self.emb(input_ids) + self.pos(pos)
+
+        for layer in self.layers:
+            h = layer(h)
+
+        return self.ln(h)
+
+
+# ============================================================
+# Causal LM Head
+# ============================================================
+
+class GCLMForCausalLM(PreTrainedModel):
+    config_class = None
+    base_model_prefix = "gclm"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.gclm = GCLMModel(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids,
+        labels=None,
+        **kwargs
+    ):
+        hidden = self.gclm(input_ids)
+        logits = self.lm_head(hidden)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits
+        )