Upload ByteETM-Korean (HF inference compatible)

modeling_byteetm.py

@@ -1,5 +1,176 @@
 from transformers import PreTrainedModel, PretrainedConfig
 import torch.nn as nn, torch.nn.functional as F, torch
+import math, random, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+
+# ---------- 4. 모델 정의 ----------
+# === GeneratingSeries 기반 보조 모듈 ===
+class MomentumEncoder(nn.Module):
+    """토큰 임베딩 간의 차분을 포함한 동적 인코딩"""
+    def __init__(self, dim):
+        super().__init__()
+        self.linear = nn.Linear(dim * 2, dim)
+        self.norm = nn.LayerNorm(dim)
+        self.act = nn.Tanh()
+    def forward(self, x):  # [B,T,C]
+        diff = F.pad(x[:, 1:] - x[:, :-1], (0,0,1,0))
+        return self.act(self.norm(self.linear(torch.cat([x, diff], dim=-1))))
+
+class GFLayer(nn.Module):
+    """지수 감쇠 기반의 생성함수 확장"""
+    def __init__(self, dim, max_order=6, tau_scale=0.01):
+        super().__init__()
+        self.coeff = nn.Parameter(torch.randn(dim, max_order + 1) * 0.1)
+        self.tau = nn.Parameter(torch.ones(dim) * tau_scale)
+        self.max_order = max_order
+    def forward(self, x):
+        B, T, D = x.shape
+        t = torch.arange(T, device=x.device).float().view(1,T,1)
+        z = torch.exp(-t * self.tau.view(1,1,D))
+        powers = torch.stack([z**k for k in range(self.max_order+1)], dim=-1)
+        gen = torch.einsum('btdk,dk->btd', powers, self.coeff)
+        return x + gen
+
+class OrthogonalTemporalProjector(nn.Module):
+    """시퀀스 길이 방향으로 직교 기저 투영"""
+    def __init__(self, t_len, rank=8):
+        super().__init__()
+        self.U = nn.Parameter(torch.randn(t_len, rank) / math.sqrt(t_len))
+    def forward(self, x):
+        B,T,D = x.shape
+        if T != self.U.size(0):
+            U = F.interpolate(self.U.T.unsqueeze(0), size=T, mode="linear", align_corners=False).squeeze(0).T
+        else:
+            U = self.U
+        U = F.normalize(U, dim=0)
+        P = U @ U.T
+        trend = torch.einsum('btd,ts->bsd', x, P)
+        resid = x - trend
+        return 0.5*(trend + resid)
+
+# === GPT Block 확장 ===
+class GeneratingBlock(nn.Module):
+    """기존 Transformer Block + GeneratingSeries 동역학 통합"""
+    def __init__(self, n_embd, n_head, block_size, dropout=0.0, gf_order=6):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(n_embd)
+        self.ln2 = nn.LayerNorm(n_embd)
+        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
+        self.mlp = MLP(n_embd, dropout)
+        # GeneratingSeries 요소
+        self.momentum = MomentumEncoder(n_embd)
+        self.gf = GFLayer(n_embd, max_order=gf_order)
+        self.otp = OrthogonalTemporalProjector(block_size, rank=min(8, block_size//4))
+    def forward(self, x):
+        # step1: momentum encoding (local diff)
+        x = self.momentum(x)
+        # step2: attention + residual
+        x = x + self.attn(self.ln1(x))
+        # step3: generating function expansion in feature domain
+        x = self.gf(x)
+        # step4: feedforward + residual
+        x = x + self.mlp(self.ln2(x))
+        # step5: orthogonal trend projection (temporal disentangling)
+        x = self.otp(x)
+        return x
+
+# === CausalSelfAttention과 MLP는 기존과 동일 ===
+class CausalSelfAttention(nn.Module):
+    def __init__(self, n_embd, n_head, block_size, dropout=0.0):
+        super().__init__()
+        assert n_embd % n_head == 0
+        self.n_head = n_head
+        self.key = nn.Linear(n_embd, n_embd)
+        self.query = nn.Linear(n_embd, n_embd)
+        self.value = nn.Linear(n_embd, n_embd)
+        self.proj = nn.Linear(n_embd, n_embd)
+        self.attn_drop = nn.Dropout(dropout)
+        self.resid_drop = nn.Dropout(dropout)
+        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)).view(1,1,block_size,block_size))
+
+    def forward(self, x):
+        B, T, C = x.size()
+        k = self.key(x).view(B, T, self.n_head, C//self.n_head).transpose(1,2)
+        q = self.query(x).view(B, T, self.n_head, C//self.n_head).transpose(1,2)
+        v = self.value(x).view(B, T, self.n_head, C//self.n_head).transpose(1,2)
+        att = (q @ k.transpose(-2,-1)) / math.sqrt(k.size(-1))
+        att = att.masked_fill(self.mask[:,:,:T,:T]==0, float("-inf"))
+        att = F.softmax(att, dim=-1)
+        att = self.attn_drop(att)
+        y = att @ v
+        y = y.transpose(1,2).contiguous().view(B,T,C)
+        y = self.resid_drop(self.proj(y))
+        return y
+
+class MLP(nn.Module):
+    def __init__(self, n_embd, dropout=0.0):
+        super().__init__()
+        self.fc = nn.Sequential(
+            nn.Linear(n_embd, 4*n_embd),
+            nn.GELU(),
+            nn.Linear(4*n_embd, n_embd),
+            nn.Dropout(dropout),
+        )
+    def forward(self, x): return self.fc(x)
+
+class Block(nn.Module):
+    def __init__(self, n_embd, n_head, block_size, dropout=0.0):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(n_embd)
+        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
+        self.ln2 = nn.LayerNorm(n_embd)
+        self.mlp = MLP(n_embd, dropout)
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+class ByteETM(nn.Module):
+    def __init__(self, vocab_size, n_embd, n_head, n_layer, block_size, dropout=0.0):
+        super().__init__()
+        self.token_emb = nn.Embedding(vocab_size, n_embd)
+        self.pos_emb   = nn.Embedding(block_size, n_embd)
+        self.drop = nn.Dropout(dropout)
+        # self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
+        self.blocks = nn.ModuleList([GeneratingBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
+        self.ln_f = nn.LayerNorm(n_embd)
+        self.head = nn.Linear(n_embd, vocab_size, bias=False)
+        self.block_size = block_size
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+        if isinstance(m, nn.Linear) and m.bias is not None:
+            nn.init.zeros_(m.bias)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.size()
+        assert T <= self.block_size
+        pos = torch.arange(0, T, device=idx.device).unsqueeze(0)
+        x = self.token_emb(idx) + self.pos_emb(pos)
+        x = self.drop(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.ln_f(x)
+        logits = self.head(x)  # (B,T,V)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / max(temperature, 1e-8)
+            if top_k is not None:
+                v, _ = torch.topk(logits, top_k)
+                logits[logits < v[:, [-1]]] = -float("inf")
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, next_id), dim=1)
+        return idx
 
 class ByteETMConfig(PretrainedConfig):
     model_type = "byteetm"
@@ -15,7 +186,6 @@ class HFByteETM(PreTrainedModel):
     config_class = ByteETMConfig
     def __init__(self, config):
         super().__init__(config)
-        from model import ByteETM     # 네가 정의한 실제 모델
         self.model = ByteETM(
             vocab_size=config.vocab_size,
             n_embd=config.n_embd,