Publish TinyGPT checkpoint

README.md

@@ -0,0 +1,66 @@
+---
+license: apache-2.0
+tags:
+  - pytorch
+  - gpt
+  - tiny-gpt
+  - causal-lm
+---
+# tiny-gpt-2-1m
+
+This repository contains a pretrained TinyGPT checkpoint published for public use.
+This checkpoint is provided for educational and experimentation purposes.
+
+## Artifacts
+
+- `tiny_gpt_latest.pt`: training checkpoint with model and optimizer state
+- `tokenizer.model`: SentencePiece tokenizer used for training and generation
+- `config.json`: model configuration serialized from the checkpoint
+- `training_config.yaml`: training and MLflow settings used for the run
+
+## How to use
+
+Use with Transformers.
+
+Starting with `transformers >= 4.43.0`, you can run conversational inference using the `pipeline` abstraction or by leveraging the `Auto` classes with `generate()`.
+
+Make sure to update your Transformers installation via `pip install --upgrade transformers`.
+
+```python
+import torch
+import transformers
+
+model_id = "vjkhambe/tiny-gpt-2-1m"
+device = 0 if torch.cuda.is_available() else -1
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+model = transformers.AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    dtype=dtype,
+)
+tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model.generation_config.max_length = None
+model.generation_config.max_new_tokens = 64
+
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=device,
+)
+
+print(pipeline("Hey how are you doing today?"))
+```
+
+## Training details
+
+- Base package: `tiny_gpt_pretrain`
+- Model and training configuration are stored in the checkpoint and `training_config.yaml`
+- The exported checkpoint includes optimizer state for continued fine-tuning or evaluation
+
+## License
+
+Released under the Apache-2.0 license.
+
+Target repo: `vjkhambe/tiny-gpt-2-1m`

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "TinyGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_tiny_gpt.TinyGPTConfig",
+    "AutoModelForCausalLM": "modeling_tiny_gpt.TinyGPTForCausalLM"
+  },
+  "bos_token_id": 1,
+  "context_length": 256,
+  "d_ff": 512,
+  "d_model": 128,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_size": 128,
+  "intermediate_size": 512,
+  "is_decoder": true,
+  "max_position_embeddings": 256,
+  "model_type": "tiny_gpt",
+  "n_heads": 4,
+  "n_layers": 4,
+  "num_attention_heads": 4,
+  "num_hidden_layers": 4,
+  "pad_token_id": 3,
+  "tie_embeddings": false,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.12.1",
+  "unk_token_id": 0,
+  "vocab_size": 2048
+}

export_metadata.json

@@ -0,0 +1,5 @@
+{
+  "source_checkpoint": "checkpoints/tiny_gpt_latest.pt",
+  "source_tokenizer": "data/processed/tiny_bpe.model",
+  "transformers_compatible": true
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 3,
+  "transformers_version": "5.12.1",
+  "use_cache": false
+}

model.safetensors

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

modeling_tiny_gpt.py

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers import GenerationMixin, PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class TinyGPTConfig(PretrainedConfig):
+    model_type = "tiny_gpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 2048,
+        context_length: int = 256,
+        n_layers: int = 4,
+        n_heads: int = 4,
+        d_model: int = 128,
+        d_ff: int = 512,
+        dropout: float = 0.1,
+        tie_embeddings: bool = True,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        pad_token_id: int = 3,
+        unk_token_id: int = 0,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> None:
+        is_decoder = kwargs.pop("is_decoder", True)
+        tie_word_embeddings = kwargs.pop("tie_word_embeddings", tie_embeddings)
+        use_cache = kwargs.pop("use_cache", use_cache)
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            unk_token_id=unk_token_id,
+            use_cache=use_cache,
+            tie_word_embeddings=tie_word_embeddings,
+            is_decoder=is_decoder,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.context_length = context_length
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.dropout = dropout
+        self.tie_embeddings = tie_embeddings
+
+        # Standard Transformer aliases used by generation helpers.
+        self.hidden_size = d_model
+        self.intermediate_size = d_ff
+        self.num_attention_heads = n_heads
+        self.num_hidden_layers = n_layers
+        self.max_position_embeddings = context_length
+        self.head_dim = d_model // n_heads
+
+        if self.d_model % self.n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(self, config: TinyGPTConfig) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.embedding.weight
+
+    def forward(self, idx: torch.Tensor) -> torch.Tensor:
+        return self.embedding(idx)
+
+
+class PositionEmbedding(nn.Module):
+    def __init__(self, config: TinyGPTConfig) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(config.context_length, config.d_model)
+
+    def forward(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        positions = torch.arange(seq_len, device=device).unsqueeze(0)
+        return self.embedding(positions)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: TinyGPTConfig) -> None:
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.head_dim = config.d_model // config.n_heads
+
+        self.q_proj = nn.Linear(config.d_model, config.d_model)
+        self.k_proj = nn.Linear(config.d_model, config.d_model)
+        self.v_proj = nn.Linear(config.d_model, config.d_model)
+        self.out_proj = nn.Linear(config.d_model, config.d_model)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+
+        mask = torch.tril(torch.ones(config.context_length, config.context_length))
+        self.register_buffer(
+            "causal_mask",
+            mask.view(1, 1, config.context_length, config.context_length),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, d_model = x.shape
+
+        query = self._split_heads(self.q_proj(x), batch_size, seq_len)
+        key = self._split_heads(self.k_proj(x), batch_size, seq_len)
+        value = self._split_heads(self.v_proj(x), batch_size, seq_len)
+
+        scores = query @ key.transpose(-2, -1)
+        scores = scores / (self.head_dim**0.5)
+        scores = scores.masked_fill(
+            self.causal_mask[:, :, :seq_len, :seq_len] == 0,
+            float("-inf"),
+        )
+
+        attention_weights = F.softmax(scores, dim=-1)
+        attention_weights = self.attn_dropout(attention_weights)
+
+        out = attention_weights @ value
+        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
+        out = self.out_proj(out)
+        return self.resid_dropout(out)
+
+    def _split_heads(self, x: torch.Tensor, batch_size: int, seq_len: int) -> torch.Tensor:
+        x = x.view(batch_size, seq_len, self.n_heads, self.head_dim)
+        return x.transpose(1, 2)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: TinyGPTConfig) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(config.d_model, config.d_ff),
+            nn.GELU(),
+            nn.Linear(config.d_ff, config.d_model),
+            nn.Dropout(config.dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: TinyGPTConfig) -> None:
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.d_model)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.d_model)
+        self.mlp = FeedForward(config)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class TinyGPTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = TinyGPTConfig
+    base_model_prefix = "tiny_gpt"
+    main_input_name = "input_ids"
+    _tied_weights_keys = {"lm_head.weight": "token_embedding.embedding.weight"}
+
+    def __init__(self, config: TinyGPTConfig) -> None:
+        super().__init__(config)
+
+        self.token_embedding = TokenEmbedding(config)
+        self.position_embedding = PositionEmbedding(config)
+        self.dropout = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layers))
+        self.final_ln = nn.LayerNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        self.post_init()
+        self.tie_weights()
+        if getattr(self, "generation_config", None) is not None:
+            self.generation_config.use_cache = False
+            self.generation_config.cache_implementation = None
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.token_embedding.embedding
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.token_embedding.embedding = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.lm_head = new_embeddings
+
+    def tie_weights(self, *args, **kwargs) -> None:
+        del args, kwargs
+        if self.config.tie_embeddings:
+            self.lm_head.weight = self.token_embedding.weight
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            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=0.02)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        past_key_values=None,
+        use_cache: bool | None = None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]:
+        del attention_mask, past_key_values, use_cache, kwargs
+
+        batch_size, seq_len = input_ids.shape
+        if seq_len > self.config.context_length:
+            raise ValueError(
+                f"Sequence length {seq_len} exceeds context length {self.config.context_length}"
+            )
+        if labels is not None and labels.shape != input_ids.shape:
+            raise ValueError(f"labels shape {labels.shape} must match input_ids shape {input_ids.shape}")
+
+        token_embeddings = self.token_embedding(input_ids)
+        position_embeddings = self.position_embedding(seq_len, input_ids.device)
+        x = token_embeddings + position_embeddings
+        x = self.dropout(x)
+
+        for block in self.blocks:
+            x = block(x)
+
+        x = self.final_ln(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            if seq_len < 2:
+                raise ValueError("Need at least 2 tokens to compute causal LM loss")
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.reshape(-1, self.config.vocab_size),
+                shift_labels.reshape(-1),
+                ignore_index=-100,
+            )
+
+        if not return_dict:
+            output = (logits,)
+            if loss is not None:
+                output = (loss,) + output
+            return output
+
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.Tensor,
+        past_key_values=None,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor | None]:
+        del past_key_values
+        if input_ids.shape[1] > self.config.context_length:
+            input_ids = input_ids[:, -self.config.context_length :]
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, -self.config.context_length :]
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "use_cache": False,
+            **kwargs,
+        }
+
+
+TinyGPTConfig.register_for_auto_class()
+TinyGPTForCausalLM.register_for_auto_class("AutoModelForCausalLM")

tiny_gpt_latest.pt

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

tokenization_tiny_gpt.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+
+import sentencepiece as spm
+from transformers import PreTrainedTokenizer
+
+
+class TinyGPTTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "tokenizer.model"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: str,
+        unk_token: str = "<unk>",
+        bos_token: str = "<s>",
+        eos_token: str = "</s>",
+        pad_token: str = "<pad>",
+        **kwargs,
+    ) -> None:
+        self.vocab_file = vocab_file
+        self.sp_model = spm.SentencePieceProcessor(model_file=vocab_file)
+        super().__init__(
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return self.sp_model.get_piece_size()
+
+    def get_vocab(self) -> dict[str, int]:
+        vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text: str) -> list[str]:
+        return list(self.sp_model.encode(text, out_type=str))
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return int(self.sp_model.piece_to_id(token))
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.sp_model.id_to_piece(int(index))
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return self.sp_model.decode_pieces(tokens)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        if token_ids_1 is None:
+            return token_ids_0
+        return token_ids_0 + token_ids_1
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+        already_has_special_tokens: bool = False,
+    ) -> list[int]:
+        if already_has_special_tokens:
+            return [0] * (len(token_ids_0) + (len(token_ids_1) if token_ids_1 else 0))
+        return [0] * (len(token_ids_0) + (len(token_ids_1) if token_ids_1 else 0))
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+        if not self.vocab_file:
+            raise ValueError("No SentencePiece model file to save")
+
+        save_dir = Path(save_directory)
+        save_dir.mkdir(parents=True, exist_ok=True)
+        filename = "tokenizer.model"
+        out_path = save_dir / filename
+        shutil.copy2(self.vocab_file, out_path)
+        return (str(out_path),)
+
+
+TinyGPTTokenizer.register_for_auto_class()

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,49 @@
+{
+  "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
+    },
+    "3": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_tiny_gpt.TinyGPTTokenizer",
+      null
+    ]
+  },
+  "backend": "custom",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "model_max_length": 1000000000,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TinyGPTTokenizer",
+  "unk_token": "<unk>"
+}

training_config.yaml

@@ -0,0 +1,30 @@
+training:
+  input: data/raw/input.txt
+  tokenizer: data/processed/tiny_bpe.model
+  checkpoint_dir: checkpoints
+  checkpoint_prefix: tiny_gpt
+  batch_size: 64
+  max_steps: 5000
+  learning_rate: 3e-4
+  weight_decay: 0.1
+  grad_clip: 1.0
+  eval_interval: 250
+  eval_steps: 50
+  train_split: 0.9
+  seed: 1337
+  device: auto
+  generate_tokens: 160
+  keep_checkpoints: 5
+  show_sample: false
+  show_checkpoint: false
+
+mlflow:
+  enabled: true
+  tracking_uri: sqlite:///mlflow.db
+  experiment_name: tiny-gpt-adoption
+  run_name: null
+  log_checkpoints: true
+  log_model_config: true
+  tags:
+    repo: tiny-gpt-adoption
+    purpose: pretrain-mlflow-hf