Upload folder using huggingface_hub

modeling_ivme.py

@@ -0,0 +1,130 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+class IvmeConfig(PretrainedConfig):
+    model_type = "ivme"
+    def __init__(self, vocab_size=16000, context_len=1024, tie_word_embeddings=True, **kwargs):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.vocab_size = vocab_size
+        self.context_len = context_len
+        self.tie_word_embeddings = tie_word_embeddings
+        self.hidden_dim = 384
+        self.n_layers = 10
+        self.n_heads = 6
+        self.dropout = 0.0
+        self.ffn_mult = 4.0
+        self.norm_eps = 1e-05
+        self.rope_theta = 10000.0
+        self.head_dim = 64
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    def forward(self, x):
+        pow_x = x.pow(2).mean(-1, keepdim=True)
+        return x * torch.rsqrt(pow_x + self.eps) * self.weight
+
+def precompute_rope_freqs(dim: int, max_seq_len: int, theta: float = 10000.0):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.outer(t, inv_freq)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, hidden_dim: int, n_heads: int, dropout: float = 0.0):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = hidden_dim // n_heads
+        self.wq = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.wk = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.wv = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.wo = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+    def forward(self, x, rope_freqs):
+        B, T, C = x.shape
+        q, k, v = self.wq(x), self.wk(x), self.wv(x)
+        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        
+        q_complex = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2))
+        k_complex = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2))
+        freqs = rope_freqs[:T].view(1, 1, T, -1).to(q_complex.device)
+        
+        q = torch.view_as_real(q_complex * freqs).flatten(3).to(x.dtype)
+        k = torch.view_as_real(k_complex * freqs).flatten(3).to(x.dtype)
+        v = v.to(x.dtype)
+        
+        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
+        mask = torch.full((T, T), float("-inf"), device=x.device).triu(1)
+        scores = scores + mask
+        probs = torch.softmax(scores, dim=-1).to(x.dtype)
+        probs = self.dropout(probs)
+        
+        output = torch.matmul(probs, v)
+        output = output.transpose(1, 2).contiguous().view(B, T, C)
+        return self.wo(output)
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden_dim: int, ffn_mult: float = 1.0):
+        super().__init__()
+        hidden_features = int(2 * hidden_dim * 4 / 3)
+        hidden_features = int(ffn_mult * hidden_features)
+        self.w1 = nn.Linear(hidden_dim, hidden_features, bias=False)
+        self.w2 = nn.Linear(hidden_features, hidden_dim, bias=False)
+        self.w3 = nn.Linear(hidden_dim, hidden_features, bias=False)
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
+        self.attn = CausalSelfAttention(cfg.hidden_dim, cfg.n_heads, cfg.dropout)
+        self.ffn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
+        self.ffn = SwiGLU(cfg.hidden_dim, cfg.ffn_mult)
+    def forward(self, x, rope_freqs):
+        x = x + self.attn(self.attn_norm(x), rope_freqs)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+class IvmeConversateV2HF(PreTrainedModel):
+    config_class = IvmeConfig
+    base_model_prefix = "model"
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        
+        self.tok_embed = nn.Embedding(config.vocab_size, config.hidden_dim)
+        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
+        self.final_norm = RMSNorm(config.hidden_dim, eps=config.norm_eps)
+        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
+        
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.tok_embed.weight
+            
+        rope_freqs = precompute_rope_freqs(config.hidden_dim // config.n_heads, config.context_len, config.rope_theta)
+        self.register_buffer("rope_freqs", rope_freqs, persistent=False)
+        
+        self.post_init()
+
+    def forward(self, input_ids=None, labels=None, **kwargs):
+        B, T = input_ids.shape
+        x = self.tok_embed(input_ids)
+        for block in self.blocks:
+            x = block(x, self.rope_freqs)
+        x = self.final_norm(x)
+        logits = self.lm_head(x)
+        
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-1)
+            
+        return CausalLMOutputWithPast(loss=loss, logits=logits)