Upload folder using huggingface_hub

.gitignore

@@ -123,9 +123,17 @@ temp/
 *.sqlite3
 *.db
 
+
+# =========================
+# Jupyter
+# =========================
+.gradio/
+
+
 # =========================
 # uv
 # =========================
 # Keep uv.lock tracked.
 # Do NOT add uv.lock to .gitignore.
 
+

README.md

@@ -11,16 +11,38 @@ short_description: A set of tiny Transformers, from scratch
 
 # Tiny Text Transformers
 
-A small Transformer built from scratch in PyTorch,
-with multiple NLP tasks built on the same Transformer core.
+---
+title: Tiny Transformers
+emoji: 🧠
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+app_file: app.py
+pinned: false
+short_description: A set of tiny Transformers, from scratch
+---
+
+# Tiny Text Transformers
+
+A set of small transformer models, less than 1M params for multiple NLP tasks built on the same Transformer core.
+
+Goal: Light, highly specialized and useful in a specific domain
 
 ## Tasks
 
 ### Text Generation
-Character-level autoregressive language model trained on Shakespeare.
+Character-level autoregressive language model. Shakespear corpus used for proof of concept.
 
 ### Text Classification
-Character-level Transformer classifier trained on AG News.
+News / text classifier. AG News dataset used for proof of concept.
 
 ### Structured Output
-Coming next...
+Generate structured output for a limited instruction-set problem. For instance a cli command.
+
+### Translation
+Encoder Decoder block for experimental educational purposes.
+
+### Hugging Faces
+https://huggingface.co/spaces/zeguru/tiny-transformers
+
+

app.py

@@ -12,7 +12,7 @@ from src.task.inference.classify import TextClassifier
 
 # Paths
 # DATA_PATH = Path("data/tiny_shakespear.txt")
-MODEL_PATH = Path("checkpoints/tiny-text-transformer.pt")
+MODEL_PATH = Path("checkpoints/tiny_text_generator.pt")
 
 CLASSIFIER_MODEL_PATH = Path(
     "checkpoints/news-classifier.pt"
@@ -55,12 +55,14 @@ def generate_text(
     prompt,
     max_new_tokens,
     temperature,
+    top_p,
     ):
 
     return text_generator.generate(
         prompt_text=prompt,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
+        top_p=top_p
         )
 
 # B. Classifier
@@ -96,41 +98,6 @@ def classify_text(text):
 
 
 
-
-
-# Gradio Interface
-# demo = gr.Interface(
-#     fn=generate_text,
-#     inputs=[
-#         gr.Textbox(
-#             label="Prompt",
-#             value="ROMEO:",
-#             ),
-#         gr.Slider(
-#             minimum=10,
-#             maximum=500,
-#             value=300,
-#             step=10,
-#             label="Max new tokens",
-#             ),
-#         gr.Slider(
-#             minimum=0.1,
-#             maximum=1.5,
-#             value=0.7,
-#             step=0.1,
-#             label="Temperature",
-#             ),
-#         ],
-#     outputs=gr.Textbox(
-#         label="Generated text",
-#         lines=15,
-#         ),
-#     title="Tiny Text Transformer",
-#     description=(
-#         "A tiny character-level Transformer built from scratch in PyTorch, demonstrated on text generation and AG News classification."
-#         ),
-#     )
-
 with gr.Blocks() as demo:
 
     gr.Markdown(
@@ -167,6 +134,13 @@ with gr.Blocks() as demo:
                     step=0.1,
                     label="Temperature",
                 ),
+                gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.1,
+                    label="Top P",
+                ),
             ],
             outputs=gr.Textbox(
                 label="Generated text",

checkpoints/tiny_text_generator.pt

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

logs/text_generator.txt

@@ -0,0 +1,45 @@
+ModelConfig(context_length=64, d_model=64, d_ff=256, num_layers=2, dropout=0.1)
+TrainingConfig(batch_size=32, learning_rate=0.001, num_steps=10000, eval_interval=1000, eval_batches=20, num_train_samples=None)
+Corpus length: 1,115,394 characters
+Vocabulary size: 67
+Token count: 1,115,394
+Train tokens: 1,003,854
+Val tokens:   55,770
+Test tokens:  55,770
+Train examples: 1,003,790
+Val examples:   55,706
+Test examples:  55,706
+TinyTransformerLM(
+  (input_representation): InputRepresentation(
+    (token_embedding): Embedding(
+      (embedding): Embedding(67, 64)
+    )
+    (positional_embedding): PositionalEmbedding(
+      (embedding): Embedding(64, 64)
+    )
+  )
+  (transformer): Transformer(
+    (blocks): ModuleList(
+      (0-1): 2 x TransformerBlock(
+        (attention): Attention(
+          (query): Linear(in_features=64, out_features=64, bias=False)
+          (key): Linear(in_features=64, out_features=64, bias=False)
+          (value): Linear(in_features=64, out_features=64, bias=False)
+        )
+        (feed_forward): FeedForward(
+          (net): Sequential(
+            (0): Linear(in_features=64, out_features=256, bias=True)
+            (1): GELU(approximate='none')
+            (2): Linear(in_features=256, out_features=64, bias=True)
+          )
+        )
+        (layer_norm_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True, bias=True)
+        (layer_norm_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True, bias=True)
+        (dropout): Dropout(p=0.1, inplace=False)
+      )
+    )
+  )
+  (head): Linear(in_features=64, out_features=67, bias=True)
+)
+Test loss: 1.7971
+Parameters: 104,003

scripts/generate_text.py

@@ -8,9 +8,9 @@ from src.task.inference.inference import TextGenerator
 
 # Runtime configs
 
-checkpoint = "checkpoints/tiny-text-transformer.pt"
+checkpoint = "checkpoints/tiny_text_generator.pt"
 
-TEXT_TOKENIZER_PATH = Path("checkpoints/shakespeare-tokenizer.json")
+TEXT_TOKENIZER_PATH = Path("checkpoints/shakespeare_tokenizer.json")
 tokenizer = CharacterTokenizer.load(
     TEXT_TOKENIZER_PATH
     )
@@ -41,6 +41,7 @@ output = text_generator.generate(
     prompt_text="ROMEO:",
     max_new_tokens=300,
     temperature=0.7,
+    top_p=0.8
 )
 
 print(output)

scripts/train_generator.py

@@ -10,28 +10,37 @@ from src.task.config import ModelConfig,TrainingConfig
 
 from src.task.model import TinyTransformerLM
 from src.task.training.autoregressive import train, evaluate
-
-
+from src.task.training.util import TrainingLogger, count_parameters
 
 def separator(title: str) -> None:
     print("\n" + "=" * 60)
     print(title)
     print("=" * 60)
 
+# Take stok
+logger = TrainingLogger(
+    "logs/text_generator.txt"
+    )
+
+
 # Runtime configs
 corpus = "data/tiny_shakespear.txt"
-checkpoint = "txt-transformer.pt"
+checkpoint = "tiny_text_generator.pt"
 
 # Configs
 separator("Configuration")
 
 torch.manual_seed(42)
 
+# Generator
 model_config = ModelConfig()
 training_config = TrainingConfig()
 
-print(model_config)
-print(training_config)
+training_config.num_steps = 10000 
+training_config.eval_interval = 1000
+
+logger.log(model_config)
+logger.log(training_config)
 
 
 # Load the corpus
@@ -43,7 +52,7 @@ text = data_path.read_text(
     encoding="utf-8",
 )
 
-print(f"Corpus length: {len(text):,} characters")
+logger.log(f"Corpus length: {len(text):,} characters")
 
 
 # Tokens
@@ -59,8 +68,8 @@ tokens = torch.tensor(
     dtype=torch.long,
 )
 
-print(f"Vocabulary size: {tokenizer.vocab_size}")
-print(f"Token count: {len(tokens):,}")
+logger.log(f"Vocabulary size: {tokenizer.vocab_size}")
+logger.log(f"Token count: {len(tokens):,}")
 
 
 # Train/Val/Test split
@@ -71,9 +80,9 @@ train_tokens, val_tokens, test_tokens = split_tokens(
     tokens,
     )
 
-print(f"Train tokens: {len(train_tokens):,}")
-print(f"Val tokens:   {len(val_tokens):,}")
-print(f"Test tokens:  {len(test_tokens):,}")
+logger.log(f"Train tokens: {len(train_tokens):,}")
+logger.log(f"Val tokens:   {len(val_tokens):,}")
+logger.log(f"Test tokens:  {len(test_tokens):,}")
 
 
 # Datasets
@@ -94,9 +103,9 @@ test_dataset = AutoRegressiveDataset(
     context_length=model_config.context_length,
     )
 
-print(f"Train examples: {len(train_dataset):,}")
-print(f"Val examples:   {len(val_dataset):,}")
-print(f"Test examples:  {len(test_dataset):,}")
+logger.log(f"Train examples: {len(train_dataset):,}")
+logger.log(f"Val examples:   {len(val_dataset):,}")
+logger.log(f"Test examples:  {len(test_dataset):,}")
 
 
 # Models
@@ -107,7 +116,7 @@ model = TinyTransformerLM(
     vocab_size=tokenizer.vocab_size,
     )
 
-print(model)
+logger.log(model)
 
 
 # Train........
@@ -130,8 +139,11 @@ test_loss = evaluate(
     training_config=training_config,
 )
 
-print(f"Test loss: {test_loss:.4f}")
+logger.log(f"Test loss: {test_loss:.4f}")
+
+
 
+logger.log(f"Parameters: {count_parameters(model):,}")
 
 # --------------------------------------------------
 # Save checkpoint

src/infra/generation.py

@@ -8,6 +8,7 @@ def prompt(
     context_length,
     max_new_tokens,
     temperature=1.0,
+    top_p=None
 ):
 
     if temperature <= 0:
@@ -29,14 +30,17 @@ def prompt(
 
         next_token_logits = (
             next_token_logits / temperature
-        )
+            )
 
-        probabilities = torch.softmax(
-            next_token_logits,
-            dim=-1,
-        )
+        if top_p is not None:
+            probabilities = get_top_p(top_p, next_token_logits)
+        else:
+            probabilities = torch.softmax(
+                next_token_logits,
+                dim=-1,
+                )
 
-        next_token = torch.multinomial(
+        next_token = torch.multinomial(     #watermark my be done here 
             probabilities,
             num_samples=1,
         )
@@ -46,4 +50,48 @@ def prompt(
             dim=1,
         )
 
-    return tokens
+    return tokens
+
+
+def get_top_p(top_p,next_logits):
+
+    probabilities = torch.softmax(
+        next_logits,
+        dim=-1,
+    )
+
+    sorted_probabilities, sorted_indices = torch.sort(
+        probabilities,
+        descending=True,
+    )
+
+    cumulative_probabilities = torch.cumsum(
+        sorted_probabilities,
+        dim=-1,
+    )
+
+    remove = cumulative_probabilities > top_p
+
+    # Keep the first token that crosses the threshold.
+    remove[..., 1:] = remove[..., :-1].clone()
+    remove[..., 0] = False
+
+    sorted_probabilities = sorted_probabilities.masked_fill(
+        remove,
+        0.0,
+        )
+
+    probabilities = torch.zeros_like(probabilities)
+
+    probabilities.scatter_(
+        -1,
+        sorted_indices,
+        sorted_probabilities,
+        )
+
+    probabilities = probabilities / probabilities.sum(
+        dim=-1,
+        keepdim=True,
+        )
+
+    return probabilities

src/task/inference/inference.py

@@ -35,6 +35,7 @@ class TextGenerator:
         prompt_text: str,
         max_new_tokens=300,
         temperature=0.7,
+        top_p=None,
     ):
 
         prompt_tokens = torch.tensor(
@@ -49,6 +50,7 @@ class TextGenerator:
             context_length=self.context_length,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
+            top_p=top_p,
         )
 
         return self.tokenizer.decode(

src/task/training/util.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+
+
+class TrainingLogger:
+
+    def __init__(self, path):
+        self.path = Path(path)
+        self.path.parent.mkdir(
+            parents=True,
+            exist_ok=True,
+        )
+
+    def log(self, message):
+        message = str(message)
+        print(message)
+
+        with self.path.open(
+            "a",
+            encoding="utf-8",
+        ) as file:
+            file.write(message + "\n")
+
+
+def count_parameters(model):
+    return sum(
+        parameter.numel()
+        for parameter in model.parameters()
+        if parameter.requires_grad
+    )
+
+    print(
+        f"Parameters: "
+        f"{count_parameters(model):,}"
+        )