Upload folder using huggingface_hub

config.json

@@ -0,0 +1,15 @@
+{
+  "architectures": [
+    "CustomTransformerForCausalLM"
+  ],
+  "d_model": 512,
+  "dim_feedforward": 2048,
+  "dropout": 0.1,
+  "group_size": 16,
+  "model_type": "custom_transformer",
+  "n_heads": 8,
+  "num_layers": 6,
+  "torch_dtype": "float32",
+  "transformers_version": "4.42.4",
+  "vocab_size": 32000
+}

llama_interpretation.py

@@ -0,0 +1,188 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch.optim as optim
+from transformers import AutoModelForCausalLM
+from transformers.modeling_utils import PreTrainedModel
+from transformers.configuration_utils import PretrainedConfig
+# Update the DecoderLayer to use the grouped MultiHeadAttention
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model, n_heads, dim_feedforward, dropout=0.1, group_size=16):
+        super(DecoderLayer, self).__init__()
+        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout, group_size)
+        self.feed_forward = PositionwiseFeedForward(d_model, dim_feedforward, dropout)
+        self.layer_norm1 = nn.LayerNorm(d_model)
+        self.layer_norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # Self-Attention Mechanism (SA)
+        norm_x = self.layer_norm1(x)
+        x = x + self.dropout(self.self_attn(norm_x, norm_x, norm_x))
+        # Feed-Forward Network (FFN)
+        norm_x = self.layer_norm2(x)
+        x = x + self.dropout(self.feed_forward(norm_x))
+        return x
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model, n_heads, dropout=0.1, group_size=16):
+        super(MultiHeadAttention, self).__init__()
+        self.query_linear = nn.Linear(d_model, d_model)
+        self.key_linear = nn.Linear(d_model, d_model)
+        self.value_linear = nn.Linear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.n_heads = n_heads
+        self.d_model = d_model
+        self.group_size = group_size
+
+    def forward(self, query, key, value):
+        # Compute attention scores
+        query = self.query_linear(query)
+        key = self.key_linear(key)
+        value = self.value_linear(value)
+        
+        # Split the input sequences into groups
+        query_groups = query.chunk(self.group_size, dim=1)
+        key_groups = key.chunk(self.group_size, dim=1)
+        value_groups = value.chunk(self.group_size, dim=1)
+
+        attention_scores = []
+        for q, k, v in zip(query_groups, key_groups, value_groups):
+            scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.d_model)
+            scores = F.softmax(scores, dim=-1)
+            scores = self.dropout(scores)
+            attention_scores.append(torch.matmul(scores, v))
+
+        # Concatenate the outputs from all groups
+        output = torch.cat(attention_scores, dim=1)
+        return output
+
+class PositionwiseFeedForward(nn.Module):
+    def __init__(self, d_model, dim_feedforward, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+    def forward(self, x):
+        x = F.relu(self.linear1(x))
+        x = self.dropout(x)
+        x = self.linear2(x)
+        return x
+
+# Update the Decoder class to use the grouped MultiHeadAttention
+class Decoder(nn.Module):
+    def __init__(self, num_layers, d_model, n_heads, dim_feedforward, dropout=0.1, group_size=16):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList([
+            DecoderLayer(d_model, n_heads, dim_feedforward, dropout, group_size) 
+            for _ in range(num_layers)
+        ])
+        self.layer_norm = nn.LayerNorm(d_model)
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        x = self.layer_norm(x)
+        return x
+
+class Embeddings(nn.Module):
+    def __init__(self, d_model, vocab_size):
+        super(Embeddings, self).__init__()
+        self.lut = nn.Embedding(vocab_size, d_model)
+        self.d_model = d_model
+
+    def forward(self, x):
+        return self.lut(x) * math.sqrt(self.d_model)
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + self.pe[:x.size(0), :]
+        return self.dropout(x)
+class RMSNorm(nn.Module):
+    def __init__(self, dim, epsilon=1e-6, scale=True):
+        super(RMSNorm, self).__init__()
+        self.epsilon = epsilon
+        self.scale = scale
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        rms = torch.sqrt(torch.mean(torch.square(x), dim=-1, keepdim=True))
+        if self.scale:
+            weight = self.weight / (rms + self.epsilon)
+            return weight * x
+        else:
+            return x / (rms + self.epsilon)
+class TransformerDecoder(nn.Module):
+    def __init__(self, num_layers, d_model, n_heads, dim_feedforward, dropout=0.1, vocab_size=10000, group_size=16):
+        super(TransformerDecoder, self).__init__()
+        self.embeddings = Embeddings(d_model, vocab_size)
+        self.positional_encoding = PositionalEncoding(d_model, dropout)
+        self.decoder = Decoder(num_layers, d_model, n_heads, dim_feedforward, dropout)
+        self.rms_norm = RMSNorm(d_model)
+        self.group_size = group_size
+
+    def forward(self, x):
+        x = self.embeddings(x)
+        x = self.positional_encoding(x)
+        x = self.decoder(x)
+        x = self.rms_norm(x)
+        return x
+class TransformerDecoderLM(nn.Module):
+    def __init__(self, num_layers, d_model, n_heads, dim_feedforward, dropout=0.1, vocab_size=10000, group_size=16):
+        super(TransformerDecoderLM, self).__init__()
+        self.transformer = TransformerDecoder(num_layers, d_model, n_heads, dim_feedforward, dropout, vocab_size, group_size)
+        self.lm_head = nn.Linear(d_model, vocab_size)
+
+    def forward(self, input_ids):
+        transformer_output = self.transformer(input_ids)
+        lm_logits = self.lm_head(transformer_output)
+        return lm_logits
+class CustomConfig(PretrainedConfig):
+    model_type = "custom_transformer"
+    def __init__(self, num_layers=6, d_model=512, n_heads=8, dim_feedforward=2048, dropout=0.1, vocab_size=10000, group_size=16, **kwargs):
+        self.num_layers = num_layers
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.dim_feedforward = dim_feedforward
+        self.dropout = dropout
+        self.vocab_size = vocab_size
+        self.group_size = group_size
+        super().__init__(**kwargs)
+
+class CustomTransformerForCausalLM(PreTrainedModel):
+    config_class = CustomConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = TransformerDecoderLM(
+            num_layers=config.num_layers,
+            d_model=config.d_model,
+            n_heads=config.n_heads,
+            dim_feedforward=config.dim_feedforward,
+            dropout=config.dropout,
+            vocab_size=config.vocab_size,
+            group_size=config.group_size
+        )
+
+    def forward(self, input_ids, labels=None):
+        logits = self.transformer(input_ids)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+        
+        return {"loss": loss, "logits": logits}
+

model.safetensors

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

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": true,
+  "model_max_length": 32768,
+  "pad_token": "<unk>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}