Forgive me

README.md

@@ -6,7 +6,7 @@ This GitHub repository presents our research on Super Tiny Language Models (STLM
 
 Our mission is to enhance the accessibility and practicality of high-performing language models across a wide range of applications by drastically reducing their computational and energy demands. We believe that our approach has the potential to contribute to more sustainable and inclusive AI development.
 
-For a comprehensive understanding of our research methodology and initial findings, we strongly encourage you to read our paper: [Super Tiny Language Models](https://arxiv.org/abs/2405.14159)
+For a comprehensive understanding of our research methodology and initial findings, we strongly encourage you to read our paper: [Super Tiny Language Models](https://arxiv.org/abs/2405.14159). If you are new to language models we provide some reasources to read up on them in the [STARTER.md](STARTER.md) file.
 
 Please note that this repository is an evolving work in progress, reflecting the ongoing nature of our research. It is subject to frequent updates and improvements as we continue to explore and refine our work on STLMs. We welcome the community's engagement with our work, value your feedback, and appreciate any contributions to this challenging but promising endeavor.
 

STARTER.md

@@ -0,0 +1,34 @@
+# Resources for Getting Started with Language Models
+## Important Papers
+While academic papers can be difficult to read, don't worry about trying to understand everything. Oftentimes the fine details are not relevant, so try to focus on understanding the main ideas. Here are some papers that are important to understand the development of language models:
+- [the original transformer paper](https://arxiv.org/abs/1706.03762): This paper is an *encoder-decoder* architecture, which is not popular these days. Main things to try to understand from the paper are the introduction of the attention mechanism.
+- [BERT](https://arxiv.org/abs/1810.04805): This paper introduces a 'masked language modelling objective' - reconstructing a partially obscured version of the original text. It is an *encoder* only model
+- [GPT-2](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf): An example of a *decoder* only model. It uses a causal language modelling objective - predicting the next word in a sentence. More importantly, it shows that just scaling up the model size/data can lead to impressive zero-shot performance on a variety of tasks
+- [GPT-3](https://arxiv.org/abs/2005.14165): This paper introduces a few-shot learning objective - predicting the next word in a sentence given a few examples of the task.
+
+## Blogs
+- [The Illustrated Transformer](http://jalammar.github.io/illustrated-transformer/): A great blog post that explains the transformer architecture in a very visual way.
+- [Illustrated GPT-2](http://jalammar.github.io/illustrated-gpt2/): Another great blog post that explains the GPT-2 architecture in a very visual way.
+
+## Huggingface
+Huggingface is the de-facto standard for working with and sharing language models. It's great for downstream tasks using pre-trained models, and they provide helpful tools for understanding the concepts behind language models as well as tools for finetuning, a massive collection of datasets, and a large collection of pre-trained models. Here are some resources to get started with Huggingface:
+- [Huggingface Transformers Documentation](https://huggingface.co/transformers/): The official documentation for the Huggingface Transformers library.
+- [Huggingface Datasets Documentation](https://huggingface.co/docs/datasets/): The official documentation for the Huggingface Datasets library.
+- [Huggingface Model Hub](https://huggingface.co/models): The official model hub for Huggingface. You can find a large collection of pre-trained models here.
+
+## Videos
+While we can't vouch for them [Stanford](https://www.youtube.com/playlist?list=PLoROMvodv4rMFqRtEuo6SGjY4XbRIVRd4) and [MIT]() have some great free lectures on deep learning and NLP.
+
+## Repositories
+We highly recommend trying to train a language model from scratch to better understand the concepts. Here are some repositories that can help you get started:
+- [MinGPT](https://github.com/karpathy/minGPT) - Very simple PyTorch implementation of GPT-2, also the basis of this repository
+
+## Key Concepts
+- **Tokenization**: The process of breaking up text into smaller pieces, usually words or subwords. These have integer representations that can be used as input to a model.
+- **Embeddings**: A way to represent words as vectors. These vectors are learned during training and are used as input to the model. Additionally the final outputs of the model also act as embeddings.
+- **Logits**: The raw output of a model before it is converted to probabilities.
+- **Feed Forward Neural Network**: A neural network with (typically) multiple, fully connected layers that each perform a linear transformation followed by a non-linear activation function. These are applied to each token independently.
+- **Attention**: The attention mechanism learns to update the representation of each token based on other tokens. A score is computed that determines how much influence a given token has on another token. Then the "value vectors" of the tokens are combined based on these scores to determine the update
+- **Heads**: Attention heads basically divide the embedding space into multiple subspaces and then apply the attention mechanism to each subspace. This allows the model to learn different types of relationships between tokens.
+- **Transformers**: A model that has interwoven attention and feed forward neural networks. The transformer architecture is composed of multiple layers of these blocks. The transformer architecture is the basis for many modern language models.
+- **Weight Tying**: This is a technique where two layers of a neural network share the same weights -- this forces them to always compute the same function.

models/components/tokenizers/utils.py

@@ -6,7 +6,7 @@
 import unicodedata
 from collections import Counter
 
-import hydra  # to get the absolute path to the tokenizer
+from hydra.utils import to_absolute_path
 
 
 def get_tokenizer_path(tokenizer_type, vocab_size, dataset_name):
@@ -16,7 +16,7 @@ def get_tokenizer_path(tokenizer_type, vocab_size, dataset_name):
     tokenizer_folder = os.path.join(
         "models", "components", "tokenizers", "tokenizer_models"
     )
-    tokenizer_folder = hydra.utils.to_absolute_path(tokenizer_folder)
+    tokenizer_folder = to_absolute_path(tokenizer_folder)
     tokenizer_full_path = os.path.join(
         tokenizer_folder, f"{tokenizer_type}_{dataset_name}_{vocab_size}.model"
     )

models/experimental/byte_level/byte_model_shell.py

@@ -6,15 +6,15 @@
 import torch
 
 from models import core_models, embedding_models, model_heads
-from models.model_shell import ModelShell 
-
+from models.model_shell import ModelShell
 
 
 class ByteModelShell(ModelShell):
     """
     Slight deviation from the standard Model Shell to
     allow for a re-constructive auxiliary loss to the input.
     """
+
     def __init__(
         self,
         embedding_model: embedding_models.EmbedderInterface,
@@ -37,12 +37,10 @@ def forward(self, token_ids):
         # to get B, S, H (with pos encoding if necessary)
         x = self.embedding_model(token_ids)
 
-        # calculate the reconstruction loss 
+        # calculate the reconstruction loss
         logits = self.model_head(x)[0]
         loss = torch.nn.functional.cross_entropy(
-            logits.view(-1, logits.size(-1)), 
-            token_ids.view(-1), 
-            ignore_index=257
+            logits.view(-1, logits.size(-1)), token_ids.view(-1), ignore_index=257
         )
 
         # pass the embeddings through the core model
@@ -51,6 +49,4 @@ def forward(self, token_ids):
         # pass the core model output through the model head
         x = self.model_head(x)[0]
 
-        return x, loss 
-
-
+        return x, loss

models/experimental/hugging_face.py

@@ -10,10 +10,11 @@
 from trainers.base_trainer import BaseTrainer
 from trainers.dataloader import BaseDataloader
 
-def build_model(model_cfg):
-    '''
+
+def build_hf_model(model_cfg):
+    """
     Helper function to build a model from the huggingface model hub.
-    '''
+    """
     ## get the model string
     model_str = model_cfg["model_string"]
 
@@ -79,7 +80,7 @@ def __init__(self, model_cfg):
         self.model_cfg = model_cfg
         model_string = model_cfg["model_string"]
         self.tokenizer = HFTokenizerWrapper(model_string)
-        self.embeddings = build_model(model_cfg).get_input_embeddings()
+        self.embeddings = build_hf_model(model_cfg).get_input_embeddings()
 
     def decode(self, token_ids):
         """
@@ -128,7 +129,7 @@ class HFTransformerCore(torch.nn.Module):
 
     def __init__(self, model_cfg):
         super().__init__()
-        self.model = build_model(model_cfg = model_cfg)
+        self.model = build_hf_model(model_cfg=model_cfg)
 
         ## freeze the parameters
         print("Note: Freezing the parameters of the hf_core model.")
@@ -140,13 +141,14 @@ def forward(self, x):
         Calls the huggingface model in question, and returns the last hidden state.
         """
         ## get the hidden states
-        hidden_states = self.model(inputs_embeds = x, output_hidden_states = True).hidden_states
+        hidden_states = self.model(
+            inputs_embeds=x, output_hidden_states=True
+        ).hidden_states
 
         ## return the last hidden state
         if isinstance(hidden_states, tuple):
             return hidden_states[-1]
 
-
 
 class HFLMHead(torch.nn.Module):
     """
@@ -155,8 +157,8 @@ class HFLMHead(torch.nn.Module):
 
     def __init__(self, model_cfg):
         super().__init__()
-        self.lm_head = build_model(model_cfg = model_cfg).get_output_embeddings()
-    
+        self.lm_head = build_hf_model(model_cfg=model_cfg).get_output_embeddings()
+
     def forward(self, x):
         """
         Passes the input through the language model head to get logits.

models/experimental/next_thought/core_models.py

@@ -1,15 +1,16 @@
 """
 The core next-thought model.
 """
-import torch 
 
+import torch
 
 
 class BaselineCoreModel(torch.nn.Module):
     """
-    An extremely simplistic core model for 
+    An extremely simplistic core model for
     next thought prediction.
     """
+
     def __init__(self, model_cfg):
         super().__init__()
 
@@ -34,12 +35,13 @@ def forward(self, x):
             x: torch.tensor(B, S, H)
         """
         return self.model(x)
-    
+
 
 class Conv1dCoreModel(torch.nn.Module):
     """
     A core model for next thought prediction using Conv1d layers.
     """
+
     def __init__(self, model_cfg):
         super().__init__()
 
@@ -48,9 +50,6 @@ def __init__(self, model_cfg):
         self.conv2 = torch.nn.Linear(300, 300)
         self.conv3 = torch.nn.Linear(3, 3)
 
-
-
-
     def forward(self, x):
         """
         Pass an input through the model
@@ -73,4 +72,4 @@ def forward(self, x):
         x = x.view(x.size(0), 1600, 3)
         x = self.conv3(x)
         x = x.view(x.size(0), 4800)
-        return x
+        return x

models/experimental/next_thought/embedding_models.py

@@ -1,30 +1,29 @@
 """
 The Embedding model for a VAE style sequence to sequence model.
 """
-import torch 
 
-from models.embedding_models import GenericEmbedder
-from models.components.layers.transformer_blocks import GenericTransformerBlock
+import torch
 
+from models.components.layers.transformer_blocks import GenericTransformerBlock
 from models.components.positional_encoding import build_positional_encodings
 from models.components.tokenizers import build_tokenizer
-
+from models.embedding_models import EmbedderInterface
 
 # import local components
 from models.experimental.next_thought.layers import AttentionPoolingRemoval
 
 
-
-class HierarchicalEncoder(GenericEmbedder):
+class HierarchicalEncoder(EmbedderInterface):
     """
     Accepts an arbitrary length sequence as input,
     uses the QK^T matrix to, at every layer,
-    pick the top n-percent of nodes to pool into 
-    a single token (the one paying most attention 
+    pick the top n-percent of nodes to pool into
+    a single token (the one paying most attention
     to the other should be pooled into the other token).
     """
+
     def __init__(self, model_cfg):
-        super().__init__(model_cfg=model_cfg)
+        super().__init__()
         # build the tokenizer
         self.tokenizer = build_tokenizer(
             tokenizer_type=model_cfg["embedder"]["tokenizer_type"],
@@ -41,7 +40,6 @@ def __init__(self, model_cfg):
         # build the positional encodings
         self.positional_encodings = build_positional_encodings(model_cfg=model_cfg)
 
-
         self.standard_transformer = torch.nn.ModuleList(
             [
                 GenericTransformerBlock(
@@ -56,25 +54,23 @@ def __init__(self, model_cfg):
 
         self.pooling_transformer = torch.nn.ModuleList(
             [
-
                 AttentionPoolingRemoval(
                     hidden_size_in=model_cfg["embedder"]["pooling_dims"][i],
-                    hidden_size_out=model_cfg["embedder"]["pooling_dims"][i+1],
+                    hidden_size_out=model_cfg["embedder"]["pooling_dims"][i + 1],
                     num_attention_heads=12,
                     pct_pool_per_layer=model_cfg["embedder"]["pct_pool_per_layer"][i],
-                ) for i in range(len(model_cfg["embedder"]["pooling_dims"]) - 1)
+                )
+                for i in range(len(model_cfg["embedder"]["pooling_dims"]) - 1)
             ]
         )
 
-
     def forward(self, token_ids):
-        # embed the input 
+        # embed the input
         x = self.embedding(token_ids)
 
-        # apply positional encoding 
+        # apply positional encoding
         x = x + self.positional_encoding(x)
 
-
         # first pass through normal attention blocks
         for layer in self.standard:
             x = layer(x)
@@ -84,4 +80,4 @@ def forward(self, token_ids):
             x = layer(x)
         # mean pool final representation
         x = x.mean(dim=-2)
-        return x
+        return x

models/model_shell.py

@@ -8,7 +8,6 @@
 from models import core_models, embedding_models, model_heads
 
 
-
 class ModelShell(torch.nn.Module):
     """
     Unify the embedding model, core model and LM head
@@ -71,7 +70,9 @@ def inference(self, model_input):
         # check if input is string
         if isinstance(model_input, str):
             # use inference function of the embedding model
-            model_input = self.embedding_model.tokenize_input(model_input, truncate=True, add_eot=False)
+            model_input = self.embedding_model.tokenize_input(
+                model_input, truncate=True, add_eot=False
+            )
         x = torch.tensor(model_input, device=self.device, dtype=torch.long).unsqueeze(0)
         x = self.embedding_model(model_input)
 
@@ -94,9 +95,16 @@ def loglikelihood(self, prefixes, continuations):
         Returns:
             ll: torch.tensor(B)
         """
-        total_strings = [f"{prefix} {cont}" for prefix, cont in zip(prefixes, continuations)]
-        input_tokens = [self.embedding_model.tokenize_input(string, truncate=True) for string in total_strings]
-        padded_batch, mask = self.embedding_model.pad_batch(input_tokens, direction="right")
+        total_strings = [
+            f"{prefix} {cont}" for prefix, cont in zip(prefixes, continuations)
+        ]
+        input_tokens = [
+            self.embedding_model.tokenize_input(string, truncate=True)
+            for string in total_strings
+        ]
+        padded_batch, mask = self.embedding_model.pad_batch(
+            input_tokens, direction="right"
+        )
         input_tensor = torch.tensor(padded_batch, device=self.device, dtype=torch.long)
         logits, _ = self.forward(input_tensor)
         logits = logits[:, :-1].reshape(-1, logits.size(-1))
@@ -105,4 +113,4 @@ def loglikelihood(self, prefixes, continuations):
         mask = mask[:, 1:].reshape(-1).to(ll.device)
         ll = ll * mask
         ll = ll.view(input_tensor.size(0), -1).sum(dim=1)
-        return -ll
+        return -ll

stlm_hf_integration/configs.py

@@ -0,0 +1,15 @@
+"""Wrapper from huggingface"""
+
+from transformers import PretrainedConfig
+
+
+class STLMConfig(PretrainedConfig):
+    model_type = "stlm"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+if __name__ == "__main__":
+    stlmconfig = STLMConfig()
+    stlmconfig.save_pretrained("stlm")