Add Sherlock Pretrained tranformer

.gitattributes

@@ -0,0 +1,2 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+model.safetensors filter=lfs diff=lfs merge=lfs -text

ReadME.md

@@ -0,0 +1,48 @@
+# SPT: A Lightweight Language Model
+
+NanoLlama is a compact language model trained on Sherlock Holmes stories.
+
+## Model Details
+
+- **Model Type**: NanoLlama (Causal Language Model)
+- **Number of Layers**: 12
+- **Hidden Size**: 512
+- **Number of Attention Heads**: 16
+- **Number of KV Heads**: 16
+- **Intermediate Size**: 2048
+- **Maximum Sequence Length**: 2048
+- **Vocabulary Size**: 97 (including special tokens)
+
+## Usage
+
+You can use this model with the Hugging Face Transformers library:
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("imdatta0/spt")
+model = AutoModelForCausalLM.from_pretrained("imdatta0/spt")
+
+# Generate text
+input_text = "Sherlock and I were "
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+output = model.generate(input_ids, max_length=50, num_return_sequences=1)
+
+generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+print(generated_text)
+```
+
+## Training
+
+This model was trained on Sherlock Holmes' stories on a single A100 with a batch size of 2 and gradient accumulation steps of 32 effective batch size of 64. It was trained on 1024 length character sequences for 10000 steps.  
+
+## Limitations
+
+- The model has a limited vocabulary of 97 tokens, which may affect its performance on certain tasks or domains.
+- The maximum sequence length is 2048 tokens, which may not be sufficient for very long text generation tasks.
+
+
+## Acknowledgements
+
+ - Thanks to Andrej Karpathy for his excellent videos on how to train GPT from scratch
+ - Sir Arthur Conan Doyle for the amazing stories :)

__init__.py

@@ -0,0 +1,6 @@
+from transformers import AutoConfig, AutoModelForCausalLM
+from .configuration_spt import SPTConfig
+from .modeling_spt import SPTForCausalLM
+
+AutoConfig.register("spt", SPTConfig)
+AutoModelForCausalLM.register(SPTConfig, SPTForCausalLM)

config.json

@@ -0,0 +1,17 @@
+{
+    "architectures": ["SPTForCausalLM"],
+    "model_type": "spt",
+    "n_layers": 12,
+    "vocab_size": 97,
+    "hidden_size": 512,
+    "n_attn_heads": 16,
+    "n_kv_heads": 16,
+    "intermediate_size": 2048,
+    "max_len": 2048,
+    "residual": true,
+    "normalise": true,
+    "bos_token_id": 95,
+    "eos_token_id": 95,
+    "pad_token_id": 95,
+    "unk_token_id": 96
+}

configuration_spt.py

@@ -0,0 +1,39 @@
+from transformers import PretrainedConfig, AutoConfig
+
+class SPTConfig(PretrainedConfig):
+    model_type = "spt"
+
+    def __init__(
+        self,
+        vocab_size=97,
+        hidden_size=512,
+        n_layers=12,
+        n_attn_heads=16,
+        n_kv_heads=16,
+        intermediate_size=2048,
+        max_len=2048,
+        residual=True,
+        normalise=True,
+        pad_token_id=95,
+        bos_token_id=95,
+        eos_token_id=95,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.n_layers = n_layers
+        self.n_attn_heads = n_attn_heads
+        self.n_kv_heads = n_kv_heads
+        self.intermediate_size = intermediate_size
+        self.max_len = max_len
+        self.residual = residual
+        self.normalise = normalise
+        
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs
+        )
+
+AutoConfig.register("spt", SPTConfig)

model.safetensors

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

modeling_spt.py

@@ -0,0 +1,143 @@
+import math
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, AutoModelForCausalLM
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from transformers.file_utils import add_start_docstrings_to_model_forward
+from .configuration_spt import SPTConfig
+
+def repeat_kv(hidden_states, repeat_times):
+    if repeat_times == 1:
+        return hidden_states
+    batch, n_kv_heads, seq_len, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:,:,None,:,:].expand(batch, n_kv_heads, repeat_times, seq_len, head_dim)
+    return hidden_states.reshape(batch, n_kv_heads*repeat_times, seq_len, head_dim)
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.head_dim = config.hidden_size // config.n_attn_heads
+        kv_size = config.n_kv_heads * self.head_dim
+        self.hidden_size = config.hidden_size
+        self.n_attn_heads = config.n_attn_heads
+        self.n_kv_heads = config.n_kv_heads
+
+        self.q = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.k = nn.Linear(config.hidden_size, kv_size, bias=False)
+        self.v = nn.Linear(config.hidden_size, kv_size, bias=False)
+
+        self.register_buffer('tril', torch.tril(torch.ones(config.max_len, config.max_len)).view(1, 1, config.max_len, config.max_len))
+
+    def forward(self, x):
+        batch_size, seq_len, hidden_dim = x.shape
+
+        q = self.q(x)
+        k = self.k(x)
+        v = self.v(x)
+
+        q = q.view(batch_size, seq_len, self.n_attn_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch_size, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch_size, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        k = repeat_kv(k, self.n_attn_heads//self.n_kv_heads)
+        v = repeat_kv(v, self.n_attn_heads//self.n_kv_heads)
+
+        attention = (q @ k.transpose(-2,-1)) * (1.0/math.sqrt(self.hidden_size))
+        attention = attention.masked_fill(self.tril[:,:,:seq_len,:seq_len]==0, float('-inf'))
+        probs = nn.functional.softmax(attention,dim=-1)
+        y = probs@v
+        y = y.transpose(1,2).contiguous().reshape(batch_size, seq_len, -1)
+
+        return y
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.act_fn = nn.GELU()
+
+    def forward(self,x):
+        up = self.up(x)
+        gate = self.gate(x)
+        return self.down(self.act_fn(up * gate))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attn = Attention(config)
+        self.mlp = MLP(config)
+        self.residual = config.residual
+        self.norm = RMSNorm(config.hidden_size) if config.normalise else nn.Identity()
+
+    def forward(self, x):
+        if self.residual:
+            x = x + self.attn(self.norm(x))
+            x = x + self.mlp(self.norm(x))
+        else:
+            x = self.attn(self.norm(x))
+            x = self.mlp(self.norm(x))
+        return x
+
+class SPTModel(PreTrainedModel):
+    config_class = SPTConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
+        self.norm = RMSNorm(config.hidden_size) if config.normalise else nn.Identity()
+
+    def forward(self, input_ids):
+        x = self.embedding(input_ids)
+        for layer in self.layers:
+            x = layer(x)
+        x = self.norm(x)
+        return x
+
+class SPTForCausalLM(PreTrainedModel):
+    config_class = SPTConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = SPTModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def forward(self, input_ids, labels=None):
+        x = self.model(input_ids)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=logits,
+            hidden_states=x,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        return past
+
+# Register the custom model
+AutoModelForCausalLM.register(SPTConfig, SPTForCausalLM)

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+    "pad_token": "#",
+    "bos_token": "#",
+    "eos_token": "#",
+    "unk_token": "[UNK]"
+}

tokenization_spt.py

@@ -0,0 +1,167 @@
+from transformers import PreTrainedTokenizer
+from typing import List, Optional
+import json
+
+class SPTTokenizer(PreTrainedTokenizer):
+    def __init__(self, vocab_file=None, **kwargs):
+        super().__init__(**kwargs)
+        self.vocab = self.load_vocab(vocab_file)
+        self.inv_vocab = {v: k for k, v in self.vocab.items()}
+        self.pad_token = self.eos_token = "#"
+        self.unk_token = "[UNK]"
+
+    @property
+    def vocab_size(self):
+        return len(self.vocab)
+
+    def get_vocab(self):
+        return dict(self.vocab)
+
+    def _tokenize(self, text):
+        return list(text)
+
+    def _convert_token_to_id(self, token):
+        return self.vocab.get(token, self.vocab.get(self.unk_token))
+
+    def _convert_id_to_token(self, index):
+        return self.inv_vocab.get(index, self.unk_token)
+
+    def convert_tokens_to_string(self, tokens):
+        return ''.join(tokens)
+
+    def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None) -> List[int]:
+        if token_ids_1 is None:
+            return token_ids_0 + [self.eos_token_id]
+        return token_ids_0 + [self.eos_token_id] + token_ids_1 + [self.eos_token_id]
+
+    def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False) -> List[int]:
+        if already_has_special_tokens:
+            return [1 if token in [self.eos_token_id] else 0 for token in token_ids_0]
+        if token_ids_1 is None:
+            return [0] * len(token_ids_0) + [1]
+        return [0] * len(token_ids_0) + [1] + [0] * len(token_ids_1) + [1]
+
+    def create_token_type_ids_from_sequences(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None) -> List[int]:
+        if token_ids_1 is None:
+            return [0] * (len(token_ids_0) + 1)
+        return [0] * (len(token_ids_0) + 1) + [1] * (len(token_ids_1) + 1)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *init_inputs, **kwargs):
+        tokenizer = super().from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)
+        return tokenizer
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        import os
+        
+        if not os.path.isdir(save_directory):
+            os.mkdir(save_directory)
+        
+        vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
+        )
+        
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self.vocab, ensure_ascii=False))
+        
+        return (vocab_file,)
+
+    def load_vocab(self, vocab_file):
+        if vocab_file is None:
+            return {'\n': 0,
+            ' ': 1,
+            '!': 2,
+            '"': 3,
+            '&': 4,
+            "'": 5,
+            '(': 6,
+            ')': 7,
+            '*': 8,
+            ',': 9,
+            '-': 10,
+            '.': 11,
+            '0': 12,
+            '1': 13,
+            '2': 14,
+            '3': 15,
+            '4': 16,
+            '5': 17,
+            '6': 18,
+            '7': 19,
+            '8': 20,
+            '9': 21,
+            ':': 22,
+            ';': 23,
+            '?': 24,
+            'A': 25,
+            'B': 26,
+            'C': 27,
+            'D': 28,
+            'E': 29,
+            'F': 30,
+            'G': 31,
+            'H': 32,
+            'I': 33,
+            'J': 34,
+            'K': 35,
+            'L': 36,
+            'M': 37,
+            'N': 38,
+            'O': 39,
+            'P': 40,
+            'Q': 41,
+            'R': 42,
+            'S': 43,
+            'T': 44,
+            'U': 45,
+            'V': 46,
+            'W': 47,
+            'X': 48,
+            'Y': 49,
+            'Z': 50,
+            '[': 51,
+            ']': 52,
+            '`': 53,
+            'a': 54,
+            'b': 55,
+            'c': 56,
+            'd': 57,
+            'e': 58,
+            'f': 59,
+            'g': 60,
+            'h': 61,
+            'i': 62,
+            'j': 63,
+            'k': 64,
+            'l': 65,
+            'm': 66,
+            'n': 67,
+            'o': 68,
+            'p': 69,
+            'q': 70,
+            'r': 71,
+            's': 72,
+            't': 73,
+            'u': 74,
+            'v': 75,
+            'w': 76,
+            'x': 77,
+            'y': 78,
+            'z': 79,
+            '£': 80,
+            '°': 81,
+            'ß': 82,
+            'à': 83,
+            'â': 84,
+            'è': 85,
+            'é': 86,
+            'ê': 87,
+            'î': 88,
+            'ñ': 89,
+            'ô': 90,
+            'ö': 91,
+            'û': 92,
+            'ü': 93}
+        else:
+            with open(vocab_file, 'r', encoding='utf-8') as f:
+                return json.load(f)

tokenizer.json

@@ -0,0 +1,36 @@
+{
+    "version": "1.0",
+    "truncation": null,
+    "padding": null,
+    "added_tokens": [
+      {
+        "id": 95,
+        "special": true,
+        "content": "#",
+        "single_word": false,
+        "lstrip": false,
+        "rstrip": false,
+        "normalized": false
+      },
+      {
+        "id": 96,
+        "special": true,
+        "content": "[UNK]",
+        "single_word": false,
+        "lstrip": false,
+        "rstrip": false,
+        "normalized": false
+      }
+    ],
+    "normalizer": null,
+    "pre_tokenizer": null,
+    "post_processor": null,
+    "decoder": null,
+    "model": {
+      "type": "BPE",
+      "vocab": {
+        "\n": 0, " ": 1, "!": 2, "\"": 3, "&": 4, "'": 5, "(": 6, ")": 7, "*": 8, ",": 9, "-": 10, ".": 11, "0": 12, "1": 13, "2": 14, "3": 15, "4": 16, "5": 17, "6": 18, "7": 19, "8": 20, "9": 21, ":": 22, ";": 23, "?": 24, "A": 25, "B": 26, "C": 27, "D": 28, "E": 29, "F": 30, "G": 31, "H": 32, "I": 33, "J": 34, "K": 35, "L": 36, "M": 37, "N": 38, "O": 39, "P": 40, "Q": 41, "R": 42, "S": 43, "T": 44, "U": 45, "V": 46, "W": 47, "X": 48, "Y": 49, "Z": 50, "[": 51, "]": 52, "`": 53, "a": 54, "b": 55, "c": 56, "d": 57, "e": 58, "f": 59, "g": 60, "h": 61, "i": 62, "j": 63, "k": 64, "l": 65, "m": 66, "n": 67, "o": 68, "p": 69, "q": 70, "r": 71, "s": 72, "t": 73, "u": 74, "v": 75, "w": 76, "x": 77, "y": 78, "z": 79, "£": 80, "°": 81, "ß": 82, "à": 83, "â": 84, "è": 85, "é": 86, "ê": 87, "î": 88, "ñ": 89, "ô": 90, "ö": 91, "û": 92, "ü": 93, "'": 94, "#": 95, "[UNK]": 96
+      },
+      "merges": []
+    }
+  }

tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+    "model_max_length": 2048,
+    "pad_token": "#",
+    "bos_token": "#",
+    "eos_token": "#",
+    "unk_token": "[UNK]",
+    "tokenizer_class": "NanoLlamaTokenizer"
+  }