Upload 10 files

Whole setup of model

__init__.py

@@ -0,0 +1,2 @@
+from .configuration_tiny import TinyConfig
+from .modeling_tiny import TinyModel, TinyForCausalLM

config.json

@@ -0,0 +1,25 @@
+{
+    "architectures": [
+        "TinyForCausalLM"
+    ],
+    "model_type": "tiny",
+    "auto_map": {
+        "AutoConfig": "configuration_tiny.TinyConfig",
+        "AutoModel": "modeling_tiny.TinyModel",
+        "AutoModelForCausalLM": "modeling_tiny.TinyForCausalLM"
+    },
+    "vocab_size": 8000,
+    "hidden_size": 384,
+    "intermediate_size": 1024,
+    "num_hidden_layers": 8,
+    "num_attention_heads": 6,
+    "max_position_embeddings": 512,
+    "rms_norm_eps": 1e-05,
+    "rope_theta": 10000.0,
+    "tie_word_embeddings": true,
+    "bos_token_id": 2,
+    "eos_token_id": 3,
+    "pad_token_id": 0,
+    "torch_dtype": "float32",
+    "transformers_version": "4.57.0"
+}

configuration_tiny.py

@@ -0,0 +1,62 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class TinyConfig(PretrainedConfig):
+    """
+    Configuration class for TinyStories Transformer.
+    """
+
+    model_type = "tiny"
+
+    def __init__(
+        self,
+        vocab_size=8000,
+        hidden_size=384,
+        intermediate_size=1024,
+        num_hidden_layers=8,
+        num_attention_heads=6,
+        max_position_embeddings=512,
+        rms_norm_eps=1e-5,
+        rope_theta=10000.0,
+        tie_word_embeddings=True,
+        bos_token_id=2,
+        eos_token_id=3,
+        pad_token_id=0,
+        initializer_range=0.02,
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+
+        self.hidden_size = hidden_size
+
+        self.intermediate_size = intermediate_size
+
+        self.num_hidden_layers = num_hidden_layers
+
+        self.num_attention_heads = num_attention_heads
+
+        self.max_position_embeddings = max_position_embeddings
+
+        self.rms_norm_eps = rms_norm_eps
+
+        self.rope_theta = rope_theta
+
+        self.initializer_range = initializer_range
+
+        self.hidden_dropout = hidden_dropout
+
+        self.attention_dropout = attention_dropout
+
+    @property
+    def head_dim(self):
+        return self.hidden_size // self.num_attention_heads

generation_config.json

@@ -0,0 +1,9 @@
+{
+    "bos_token_id": 2,
+    "eos_token_id": 3,
+    "pad_token_id": 0,
+    "do_sample": true,
+    "temperature": 0.8,
+    "top_p": 0.95,
+    "top_k": 50
+}

model.safetensors

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

modeling_tiny.py

@@ -0,0 +1,446 @@
+import math
+import sys
+from typing import Optional, Tuple, Union
+
+# Monkeypatch safetensors to handle None/missing metadata which crashes transformers
+try:
+    import safetensors
+    original_safe_open = safetensors.safe_open
+
+    class SafeOpenWrapper:
+        def __init__(self, original_obj):
+            self.original_obj = original_obj
+
+        def __enter__(self):
+            self.original_obj.__enter__()
+            return self
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            return self.original_obj.__exit__(exc_type, exc_val, exc_tb)
+
+        def metadata(self):
+            meta = self.original_obj.metadata()
+            if meta is None:
+                return {"format": "pt"}
+            return meta
+
+        def __getattr__(self, name):
+            return getattr(self.original_obj, name)
+
+    def patched_safe_open(*args, **kwargs):
+        f = original_safe_open(*args, **kwargs)
+        return SafeOpenWrapper(f)
+
+    safetensors.safe_open = patched_safe_open
+
+    if "transformers.modeling_utils" in sys.modules:
+        import transformers.modeling_utils
+        transformers.modeling_utils.safe_open = patched_safe_open
+except Exception:
+    pass
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+
+from .configuration_tiny import TinyConfig
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    assert dim % 2 == 0
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+    t = torch.arange(end)
+    freqs = torch.outer(t, freqs).float()
+    cos = torch.cos(freqs)
+    sin = torch.sin(freqs)
+    cos = torch.cat([cos, cos], dim=-1)
+    sin = torch.cat([sin, sin], dim=-1)
+    return cos, sin
+
+
+def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    T = x.shape[2]
+    cos_t = cos[:T, :].unsqueeze(0).unsqueeze(1)
+    sin_t = sin[:T, :].unsqueeze(0).unsqueeze(1)
+    return (x * cos_t) + (rotate_half(x) * sin_t)
+
+
+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: torch.Tensor) -> torch.Tensor:
+        variance = x.pow(2).mean(-1, keepdim=True)
+        return x * torch.rsqrt(variance + self.eps) * self.weight
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: TinyConfig):
+        super().__init__()
+        self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.dropout = nn.Dropout(config.hidden_dropout) if config.hidden_dropout > 0.0 else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = F.silu(self.w1(x)) * self.w2(x)
+        out = self.w3(out)
+        if self.dropout is not None:
+            out = self.dropout(out)
+        return out
+
+
+class Attention(nn.Module):
+    def __init__(self, config: TinyConfig):
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.hidden_size = config.hidden_size
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        
+        assert self.n_heads * self.head_dim == self.hidden_size
+        
+        self.wq = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.wk = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.wv = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.wo = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.dropout_p = config.attention_dropout
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        B, T, C = x.shape
+        q = self.wq(x)
+        k = self.wk(x)
+        v = 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 = apply_rotary_emb(q, cos, sin)
+        k = apply_rotary_emb(k, cos, sin)
+        
+        dropout_p = self.dropout_p if self.training else 0.0
+        
+        if attention_mask is not None:
+            if torch.all(attention_mask == 1):
+                attn_mask = None
+                is_causal = True
+            else:
+                causal_mask = torch.tril(torch.ones((T, T), dtype=torch.bool, device=x.device))
+                padding_mask = attention_mask.to(torch.bool).unsqueeze(1).unsqueeze(2) # shape: (B, 1, 1, T)
+                attn_mask = causal_mask.unsqueeze(0).unsqueeze(1) & padding_mask # shape: (B, 1, T, T)
+                is_causal = False
+        else:
+            attn_mask = None
+            is_causal = True
+            
+        out = F.scaled_dot_product_attention(
+            q, k, v, 
+            attn_mask=attn_mask, 
+            dropout_p=dropout_p, 
+            is_causal=is_causal
+        )
+        
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.wo(out)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: TinyConfig):
+        super().__init__()
+        self.attention = Attention(config)
+        self.feed_forward = FeedForward(config)
+        self.attention_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        x = x + self.attention(self.attention_norm(x), cos, sin, attention_mask)
+        x = x + self.feed_forward(self.ffn_norm(x))
+        return x
+
+
+class TinyPreTrainedModel(PreTrainedModel):
+    config_class = TinyConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["TransformerBlock"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (TinyModel, TinyForCausalLM)):
+            module.gradient_checkpointing = value
+
+
+class TinyModel(TinyPreTrainedModel):
+    def __init__(self, config: TinyConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.tok_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [TransformerBlock(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        cos, sin = precompute_freqs_cis(
+            dim=config.hidden_size // config.num_attention_heads,
+            end=config.max_position_embeddings * 2,
+            theta=config.rope_theta,
+        )
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.tok_embeddings
+
+    def set_input_embeddings(self, value):
+        self.tok_embeddings = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            hidden_states = self.tok_embeddings(input_ids)
+        elif inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            raise ValueError("You must specify either input_ids or inputs_embeds")
+
+        T = hidden_states.shape[1]
+        cos = self.cos[:T]
+        sin = self.sin[:T]
+
+        all_hidden_states = () if output_hidden_states else None
+        
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    cos,
+                    sin,
+                    attention_mask,
+                )
+            else:
+                hidden_states = layer(
+                    hidden_states,
+                    cos,
+                    sin,
+                    attention_mask,
+                )
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            past_key_values=None,
+        )
+
+
+class TinyForCausalLM(TinyPreTrainedModel):
+    _tied_weights_keys = ["tok_embeddings.weight"]
+
+    def __init__(self, config: TinyConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.tok_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [TransformerBlock(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        
+        self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        cos, sin = precompute_freqs_cis(
+            dim=config.hidden_size // config.num_attention_heads,
+            end=config.max_position_embeddings * 2,
+            theta=config.rope_theta,
+        )
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+
+        if config.tie_word_embeddings:
+            self.output.weight = self.tok_embeddings.weight
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.tok_embeddings
+
+    def set_input_embeddings(self, value):
+        self.tok_embeddings = value
+        if self.config.tie_word_embeddings:
+            self.output.weight = value.weight
+
+    def get_output_embeddings(self):
+        return self.output
+
+    def set_output_embeddings(self, new_embeddings):
+        self.output = new_embeddings
+
+    def tie_weights(self):
+        if self.config.tie_word_embeddings:
+            self.tok_embeddings.weight = self.output.weight
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            hidden_states = self.tok_embeddings(input_ids)
+        elif inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            raise ValueError("You must specify either input_ids or inputs_embeds")
+
+        T = hidden_states.shape[1]
+        cos = self.cos[:T]
+        sin = self.sin[:T]
+
+        all_hidden_states = () if output_hidden_states else None
+        
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    cos,
+                    sin,
+                    attention_mask,
+                )
+            else:
+                hidden_states = layer(
+                    hidden_states,
+                    cos,
+                    sin,
+                    attention_mask,
+                )
+
+        hidden_states = self.norm(hidden_states)
+        logits = self.output(hidden_states)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+
+        if not return_dict:
+            output = (logits,)
+            if output_hidden_states:
+                output = output + (all_hidden_states,)
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=all_hidden_states,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> dict:
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs["attention_mask"] = attention_mask
+        model_inputs["past_key_values"] = past_key_values
+        return model_inputs

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": "<pad>",
+    "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:4cffac9db68a0db7daf412a3811d5731bef8c323414e638ba0480d9980128453
+size 370810

tokenizer_config.json

@@ -0,0 +1,50 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<unk>",
+      "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
+    },
+    "3": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}