Delete modeling_ivme.py

modeling_ivme.py

@@ -1,130 +0,0 @@
-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)