Upload ETM-Korean (HF inference compatible)

config.json

@@ -0,0 +1,18 @@
+{
+  "architectures": [
+    "HFETM"
+  ],
+  "block_size": 512,
+  "dtype": "float32",
+  "is_decoder": true,
+  "model_type": "etm",
+  "n_embd": 512,
+  "n_head": 16,
+  "n_layer": 4,
+  "transformers_version": "4.57.1",
+  "vocab_size": 30000,
+  "auto_map": {
+    "AutoModelForCausalLM": "modeling_etm.HFETM",
+    "AutoConfig": "modeling_etm.ETMConfig"
+  }
+}

model.safetensors

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

modeling_etm.py

@@ -0,0 +1,347 @@
+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):
+    """다항 차분 + 게이트 통합 (길이 보존 100%)"""
+    def __init__(self, dim, max_order=3):
+        super().__init__()
+        self.max_order = max_order
+        self.proj = nn.Linear(dim * (max_order + 1), dim)
+        self.gate = nn.Linear(dim, dim)
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        # x: (B, T, D)
+        B, T, D = x.size()
+        diffs = [x]
+
+        for k in range(1, self.max_order + 1):
+            if T <= k:
+                # 길이가 너무 짧아서 차분 불가 → 전체 zero pad
+                d = torch.zeros(B, T, D, device=x.device, dtype=x.dtype)
+            else:
+                d_raw = x[:, k:] - x[:, :-k]     # (B, T-k, D)
+                pad = torch.zeros(B, k, D, device=x.device, dtype=x.dtype)
+                d = torch.cat([pad, d_raw], dim=1)   # (B, T, D)
+
+            diffs.append(d)
+
+        concat = torch.cat(diffs, dim=-1)   # (B, T, D*(max_order+1))
+        h = self.proj(concat)
+        g = torch.sigmoid(self.gate(x))
+        return self.norm(h * g + x * (1 - g))
+
+
+
+class GFLayer(nn.Module):
+    """Adaptive polynomial generating function"""
+    def __init__(self, dim, max_order=6):
+        super().__init__()
+        self.coeff = nn.Parameter(torch.randn(dim, max_order + 1) * 0.1)
+        self.alpha = nn.Parameter(torch.randn(dim) * 0.1)
+
+    def forward(self, x):
+        B, T, D = x.shape
+        t = torch.linspace(0, 1, T, device=x.device).view(1, T, 1)
+        basis = torch.stack([(t ** k) * torch.exp(-self.alpha.view(1,1,D)*t) for k in range(self.coeff.size(1))], dim=-1)
+        gen = torch.einsum("btdk,dk->btd", basis, self.coeff)
+        return x + gen
+
+
+class OrthogonalTemporalProjector(nn.Module):
+    """Adaptive rank orthogonal projection"""
+    def __init__(self, t_len, dim, rank_ratio=0.25):
+        super().__init__()
+        rank = max(4, int(rank_ratio * math.sqrt(dim)))
+        self.U = nn.Parameter(torch.randn(t_len, rank) / math.sqrt(t_len))
+
+    def forward(self, x):
+        B, T, D = x.shape
+        U = F.interpolate(self.U.T.unsqueeze(0), size=T, mode="linear", align_corners=False).squeeze(0).T
+        U = F.normalize(U, dim=0)
+        P = U @ U.T
+        trend = torch.einsum("btd,ts->bsd", x, P)
+        resid = x - trend
+        return trend + 0.5 * resid
+
+class SinusoidalPositionalEncoding(nn.Module):
+    def __init__(self, dim, max_len=2048):
+        super().__init__()
+        pe = torch.zeros(max_len, dim)
+        pos = torch.arange(0, max_len).unsqueeze(1)
+        div = torch.exp(torch.arange(0, dim, 2) * (-math.log(10000.0) / dim))
+        pe[:, 0::2] = torch.sin(pos * div)
+        pe[:, 1::2] = torch.cos(pos * div)
+        self.register_buffer("pe", pe.unsqueeze(0))
+
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1)]
+
+
+# === GPT Block 확장 ===
+class GeneratingBlock(nn.Module):
+    """기존 Transformer Block + GeneratingSeries 동역학 통합"""
+    def __init__(self, n_embd, n_head, block_size, dropout=0.0, gf_order=2):
+        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, n_embd)
+
+    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)
+
+        # RMS normalization per head
+        q = q / (q.pow(2).mean(-1, keepdim=True).sqrt() + 1e-6)
+        k = k / (k.pow(2).mean(-1, keepdim=True).sqrt() + 1e-6)
+
+        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).transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_drop(self.proj(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_enc   = SinusoidalPositionalEncoding(n_embd, max_len=block_size)
+        self.drop = nn.Dropout(dropout)
+
+        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
+        x = self.token_emb(idx)
+        x = self.pos_enc(x)          # ← 여기서 사인·코사인 위치 정보 추가
+        x = self.drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.ln_f(x)
+        logits = self.head(x)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        return logits, loss
+
+    # ====================== ByteLM 최적화 샘플러 ======================
+    @staticmethod
+    def _sample_next_token(
+        logits,              # (1, vocab_size)
+        prev_tokens,         # (1, T)
+        temperature: float = 0.7,
+        top_k: int | None = 64,
+        top_p: float | None = 0.9,
+        repetition_penalty: float = 1.1,
+        typical_p: float | None = None,
+    ):
+        """
+        Byte-level LM용 logit 후처리 + 샘플링:
+        - temperature
+        - repetition penalty
+        - top-k
+        - top-p (nucleus)
+        - optional typical sampling
+        """
+
+        # 배치 1 가정 (지금 사용 패턴 기준)
+        assert logits.size(0) == 1, "현재 샘플러는 batch=1 사용을 가정한다."
+
+        # 1) temperature scaling
+        logits = logits / max(temperature, 1e-6)
+
+        # 2) repetition penalty (이전에 나온 토큰들 확률 낮추기)
+        if repetition_penalty is not None and repetition_penalty != 1.0:
+            unique_tokens = prev_tokens.unique()
+            # 단순하게: 이전 토큰들의 logit을 나눠서 확률 감소
+            logits[:, unique_tokens] /= repetition_penalty
+
+        # 3) top-k (상위 k개만 남기기)
+        if top_k is not None and top_k > 0 and top_k < logits.size(-1):
+            v, _ = torch.topk(logits, top_k)
+            logits[logits < v[:, [-1]]] = -float("inf")
+
+        # 4) 정렬 후 top-p / typical sampling
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+        sorted_probs = F.softmax(sorted_logits, dim=-1)
+
+        # 4-1) typical sampling (선택적)
+        if typical_p is not None:
+            log_probs = torch.log(sorted_probs + 1e-12)
+            entropy = -(sorted_probs * log_probs).sum(-1, keepdim=True)
+            # https://arxiv.org/abs/2202.00666 typical sampling 구현
+            shifted_kl = torch.cumsum(sorted_probs * (entropy - log_probs), dim=-1)
+            typical_mask = shifted_kl > typical_p
+            if typical_mask.any():
+                first_idx = torch.nonzero(typical_mask[0], as_tuple=False)[0, 0]
+                sorted_logits[:, first_idx:] = -float("inf")
+                sorted_probs = F.softmax(sorted_logits, dim=-1)
+
+        # 4-2) nucleus(top-p) sampling
+        if top_p is not None and 0.0 < top_p < 1.0:
+            cumulative = torch.cumsum(sorted_probs, dim=-1)
+            # top_p를 넘는 지점부터 다 날림
+            cutoff_mask = cumulative > top_p
+            if cutoff_mask.any():
+                first_cut = torch.nonzero(cutoff_mask[0], as_tuple=False)[0, 0]
+                sorted_logits[:, first_cut:] = -float("inf")
+
+        # 5) 정렬 이전 인덱스로 복원
+        filtered_logits = torch.full_like(logits, -float("inf"))
+        filtered_logits.scatter_(1, sorted_idx, sorted_logits)
+
+        # 6) 최종 확률 분포에서 샘플링
+        probs = F.softmax(filtered_logits, dim=-1)
+
+        # ========= 안정화: 전부 NaN 또는 전부 0인 경우 대응 =========
+        if torch.isnan(probs).any() or torch.isinf(probs).any() or probs.sum() == 0:
+            # fallback: 원래 logits에서 가장 큰 토큰을 강제로 선택
+            next_id = torch.argmax(logits, dim=-1, keepdim=True)
+            return next_id
+
+        next_id = torch.multinomial(probs, num_samples=1)
+        return next_id
+
+
+    @torch.no_grad()
+    def generate(
+        self,
+        idx,
+        max_new_tokens: int = 200,
+        temperature: float = 0.7,
+        top_k: int | None = 64,
+        top_p: float | None = 0.9,
+        repetition_penalty: float = 1.1,
+        typical_p: float | None = None,
+        eos_token: int | None = None,
+    ):
+        """
+        ByteLM용 고급 generate():
+        - temperature, top_k, top_p, repetition_penalty, typical_p 지원
+        - eos_token 설정 시 해당 토큰 나오면 조기 종료
+        """
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.block_size:]          # (1, T')
+            logits, _ = self(idx_cond)                    # (1, T', V)
+            last_logits = logits[:, -1, :]                # (1, V)
+
+            next_id = self._sample_next_token(
+                last_logits,
+                prev_tokens=idx,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                typical_p=typical_p,
+            )                                             # (1, 1)
+
+            idx = torch.cat((idx, next_id), dim=1)        # (1, T+1)
+
+            if eos_token is not None and next_id.item() == eos_token:
+                break
+
+        return idx
+
+class ETMConfig(PretrainedConfig):
+    model_type = "etm"
+    def __init__(self, vocab_size=256, n_embd=512, n_head=8, n_layer=6, block_size=256, **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_head = n_head
+        self.n_layer = n_layer
+        self.block_size = block_size
+
+# 3️⃣ HF 래퍼 클래스
+class HFETM(PreTrainedModel):
+    config_class = ETMConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = ByteETM(
+            vocab_size=config.vocab_size,
+            n_embd=config.n_embd,
+            n_head=config.n_head,
+            n_layer=config.n_layer,
+            block_size=config.block_size,
+        )
+    def forward(self, input_ids, **kwargs):
+        logits, _ = self.model(input_ids)
+        return {"logits": logits}
+
+    def generate(self, *args, **kwargs):   # <── 추가
+        return self.model.generate(*args, **kwargs)

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|sep|>",
+    "<|acc|>",
+    "<|tel|>",
+    "<|rrn|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,73 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|unused0|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|unused1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "30000": {
+      "content": "<|acc|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "30001": {
+      "content": "<|tel|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "30002": {
+      "content": "<|rrn|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|sep|>",
+    "<|acc|>",
+    "<|tel|>",
+    "<|rrn|>"
+  ],
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}