Upload folder using huggingface_hub

README.md

@@ -1,17 +1,13 @@
----
-title: Tiny Transformers
-emoji: 💬
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 6.5.1
-app_file: app.py
-pinned: false
-hf_oauth: true
-hf_oauth_scopes:
-- inference-api
-license: mit
-short_description: 1M parameters and below. Specialized.
----
+Train 
 
-An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
+    ```
+    cd tiny-sentence-transformer
+    python -m src.train
+    ```
+
+Inference
+
+    ```
+    cd tiny-sentence-transformer
+    python -m src.test_inference
+    ```

app.py

@@ -1,69 +1,117 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from pathlib import Path
 
+from src.infra.tokenizer import CharacterTokenizer
+from src.task.config import ModelConfig
+from src.task.model import load_model
+from src.task.inference import TextGenerator
+
+
+# --------------------------------------------------
+# Paths
+# --------------------------------------------------
+
+DATA_PATH = Path("data/tiny_shakespear.txt")
+
+MODEL_PATH = Path(
+    "checkpoints/tiny-sentence-transformer.pt"
+)
+
+
+# --------------------------------------------------
+# Configuration
+# --------------------------------------------------
+
+config = ModelConfig()
+
+
+# --------------------------------------------------
+# Tokenizer
+# --------------------------------------------------
+
+text = DATA_PATH.read_text(
+    encoding="utf-8",
+)
+
+tokenizer = CharacterTokenizer(text)
 
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
 
-    messages = [{"role": "system", "content": system_message}]
+# --------------------------------------------------
+# Model
+# --------------------------------------------------
+
+model = load_model(
+    checkpoint_path=MODEL_PATH,
+    config=config,
+    vocab_size=tokenizer.vocab_size,
+)
+
+
+# --------------------------------------------------
+# Generator
+# --------------------------------------------------
+
+text_generator = TextGenerator(
+    model=model,
+    tokenizer=tokenizer,
+    context_length=config.context_length,
+)
 
-    messages.extend(history)
 
-    messages.append({"role": "user", "content": message})
+# --------------------------------------------------
+# Gradio function
+# --------------------------------------------------
 
-    response = ""
+def generate_text(
+    prompt,
+    max_new_tokens,
+    temperature,
+):
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
+    return text_generator.generate(
+        prompt=prompt,
+        max_new_tokens=max_new_tokens,
         temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+    )
+
+
+# --------------------------------------------------
+# 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.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
+            maximum=1.5,
+            value=0.7,
+            step=0.1,
+            label="Temperature",
         ),
     ],
+    outputs=gr.Textbox(
+        label="Generated text",
+        lines=15,
+    ),
+    title="Tiny Sentence Transformer",
+    description=(
+        "A tiny character-level transformer "
+        "trained on a Shakespeare corpus."
+    ),
 )
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
-
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

checkpoints/tiny-sentence-transformer-no-position.pt

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

checkpoints/tiny-sentence-transformer.pt

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

checkpoints/ttt-autocomplete.pt

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

pyproject.toml

@@ -0,0 +1,9 @@
+[project]
+name = "tiny-sentence-transformer"
+version = "0.1.0"
+requires-python = ">=3.14"
+dependencies = [
+    "gradio>=6.26.0",
+    "numpy>=2.5.2",
+    "torch>=2.13.0",
+]

requirements.txt

@@ -0,0 +1,2 @@
+torch
+gradio

src/generate.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+
+from src.task.config import ModelConfig
+from src.task.model import load_model
+from src.infra.tokenizer import CharacterTokenizer
+from src.task.inference import TextGenerator
+
+
+# ------------------------------------------
+# Load corpus
+# ------------------------------------------
+
+text = Path(
+    "data/tiny_shakespear.txt"
+).read_text(
+    encoding="utf-8"
+)
+
+
+# ------------------------------------------
+# Tokenizer
+# ------------------------------------------
+
+tokenizer = CharacterTokenizer(text)
+
+
+# ------------------------------------------
+# Configuration
+# ------------------------------------------
+
+config = ModelConfig()
+
+
+# ------------------------------------------
+# Load trained model
+# ------------------------------------------
+
+model = load_model(
+    checkpoint_path="checkpoints/tiny-sentence-transformer.pt",
+    config=config,
+    vocab_size=tokenizer.vocab_size,
+)
+
+
+# ------------------------------------------
+# Generator
+# ------------------------------------------
+
+text_generator = TextGenerator(
+    model=model,
+    tokenizer=tokenizer,
+    context_length=config.context_length,
+)
+
+
+# ------------------------------------------
+# Generate
+# ------------------------------------------
+
+output = text_generator.generate(
+    prompt="ROMEO:",
+    max_new_tokens=300,
+    temperature=0.7,
+)
+
+
+print(output)

src/infra/dataset.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import torch
+
+from src.infra.tokenizer import CharacterTokenizer
+
+
+class TextClassificationDataset:
+
+    def __init__(
+        self,
+        texts,
+        labels,
+        tokenizer: CharacterTokenizer,
+    ):
+        if len(texts) != len(labels):
+            raise ValueError(
+                "texts and labels must have the same length"
+            )
+
+        self.texts = texts
+
+        self.labels = torch.tensor(
+            labels,
+            dtype=torch.long,
+        )
+
+        encoded = [
+            tokenizer.encode(text)
+            for text in texts
+        ]
+
+        self.tokens, self.padding_mask = tokenizer.pad(
+            encoded
+        )
+
+        self.tokens = torch.tensor(
+            self.tokens,
+            dtype=torch.long,
+        )
+
+        self.padding_mask = torch.tensor(
+            self.padding_mask,
+            dtype=torch.bool,
+        )
+
+    def __len__(self):
+        return len(self.texts)
+
+    def __getitem__(self, index):
+
+        return (
+            self.tokens[index],
+            self.padding_mask[index],
+            self.labels[index],
+        )
+
+
+class AutoRegressiveDataset:
+    """Dataset for autoregressive language modeling.
+
+    A token sequence is converted into overlapping
+    input/target pairs.
+
+    For a context length of 4:
+
+        tokens:  A B C D E
+
+        input:   A B C D
+        target:  B C D E
+
+    Each target sequence is shifted one token
+    relative to its corresponding input sequence.
+    """
+
+    def __init__(
+        self,
+        tokens: torch.Tensor,
+        context_length: int,
+    ) -> None:
+
+        if tokens.ndim != 1:
+            raise ValueError(
+                "tokens must be a 1-dimensional tensor"
+            )
+
+        if context_length <= 0:
+            raise ValueError(
+                "context_length must be greater than 0"
+            )
+
+        if len(tokens) <= context_length:
+            raise ValueError(
+                "tokens must contain more elements than "
+                "context_length"
+            )
+
+        self.tokens = tokens
+        self.context_length = context_length
+
+    def __len__(self) -> int:
+        return len(self.tokens) - self.context_length
+
+    def __getitem__(
+        self,
+        index: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        if index < 0 or index >= len(self):
+            raise IndexError(
+                "dataset index out of range"
+            )
+
+        start = index
+        end = start + self.context_length
+
+        input_ids = self.tokens[start:end]
+        target_ids = self.tokens[start + 1:end + 1]
+
+        return input_ids, target_ids
+
+
+def split_tokens(
+    tokens: torch.Tensor,
+    train_ratio: float = 0.90,
+    val_ratio: float = 0.05,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+]:
+    """Split a token sequence into train, validation and test sets.
+
+    The sequence is split chronologically rather than randomly.
+    This prevents future tokens from leaking into the training set.
+    """
+
+    if tokens.ndim != 1:
+        raise ValueError(
+            "tokens must be a 1-dimensional tensor"
+        )
+
+    if not 0 < train_ratio < 1:
+        raise ValueError(
+            "train_ratio must be between 0 and 1"
+        )
+
+    if not 0 < val_ratio < 1:
+        raise ValueError(
+            "val_ratio must be between 0 and 1"
+        )
+
+    if train_ratio + val_ratio >= 1:
+        raise ValueError(
+            "train_ratio + val_ratio must be less than 1"
+        )
+
+    n = len(tokens)
+
+    train_end = int(train_ratio * n)
+    val_end = int((train_ratio + val_ratio) * n)
+
+    train_tokens = tokens[:train_end]
+    val_tokens = tokens[train_end:val_end]
+    test_tokens = tokens[val_end:]
+
+    return (
+        train_tokens,
+        val_tokens,
+        test_tokens,
+    )
+
+
+def get_batch(
+    dataset: AutoRegressiveDataset,
+    batch_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Sample a random batch from an autoregressive dataset."""
+
+    if batch_size <= 0:
+        raise ValueError(
+            "batch_size must be greater than 0"
+        )
+
+    indices = torch.randint(
+        low=0,
+        high=len(dataset),
+        size=(batch_size,),
+    )
+
+    inputs = []
+    targets = []
+
+    for index in indices:
+
+        input_ids, target_ids = dataset[index.item()]
+
+        inputs.append(input_ids)
+        targets.append(target_ids)
+
+    return (
+        torch.stack(inputs),
+        torch.stack(targets),
+    )

src/infra/generation.py

@@ -0,0 +1,101 @@
+import torch
+
+
+@torch.no_grad()
+def generate(
+    tokens,
+    embedding,
+    model,
+    output_projection,
+    context_length,
+    max_new_tokens,
+    temperature=1.0,
+):
+
+    if temperature <= 0:
+        raise ValueError("temperature must be greater than 0")
+
+    model.eval()
+
+    for _ in range(max_new_tokens):
+
+        context = tokens[:, -context_length:]
+
+        # ------------------------------------------
+        # Forward pass
+        # ------------------------------------------
+        x = embedding(context)
+        hidden = model(x)
+        logits = output_projection(hidden)
+
+        next_token_logits = logits[:, -1, :]
+
+        next_token_logits = (
+            next_token_logits / temperature
+        )
+
+        probabilities = torch.softmax(
+            next_token_logits,
+            dim=-1,
+        )
+
+        next_token = torch.multinomial(
+            probabilities,
+            num_samples=1,
+        )
+
+        tokens = torch.cat(
+            [tokens, next_token],
+            dim=1,
+        )
+
+    return tokens
+
+
+
+@torch.no_grad()
+def generator(
+    tokens,
+    generator,
+    context_length,
+    max_new_tokens,
+    temperature=1.0,
+):
+
+    if temperature <= 0:
+        raise ValueError("temperature must be greater than 0")
+
+    generator.eval()
+
+    for _ in range(max_new_tokens):
+
+        context = tokens[:, -context_length:]
+
+        # ------------------------------------------
+        # Forward pass
+        # ------------------------------------------
+       
+        logits = generator(context)
+
+        next_token_logits = logits[:, -1, :]
+
+        next_token_logits = (
+            next_token_logits / temperature
+        )
+
+        probabilities = torch.softmax(
+            next_token_logits,
+            dim=-1,
+        )
+
+        next_token = torch.multinomial(
+            probabilities,
+            num_samples=1,
+        )
+
+        tokens = torch.cat(
+            [tokens, next_token],
+            dim=1,
+        )
+
+    return tokens

src/infra/tokenizer.py

@@ -0,0 +1,87 @@
+class CharacterTokenizer:
+
+    PAD_TOKEN = "<PAD>"
+    UNK_TOKEN = "<UNK>"
+
+
+    def __init__(self, text):
+        characters = sorted(set(text))
+
+        self.stoi = {
+            self.PAD_TOKEN: 0,
+            self.UNK_TOKEN: 1,
+        }
+
+        for character in characters:
+            self.stoi[character] = len(self.stoi)
+
+        self.itos = {
+            index: token
+            for token, index in self.stoi.items()
+        }
+
+    @property
+    def vocab_size(self):
+        return len(self.stoi)
+
+    @property
+    def pad_token_id(self):
+        return self.stoi[self.PAD_TOKEN]
+
+    @property
+    def unk_token_id(self):
+        return self.stoi[self.UNK_TOKEN]
+
+    def encode(self, text):
+        return [
+            self.stoi.get(
+                character,
+                self.unk_token_id,
+            )
+            for character in text
+        ]
+
+    def decode(self, token_ids):
+        return "".join(
+            self.itos[token_id]
+            for token_id in token_ids
+            if token_id != self.pad_token_id
+        )
+
+    def pad(
+            self,
+            sequences,
+            max_length=None,
+            ):
+        
+        if max_length is None:
+            max_length = max(
+                len(sequence)
+                for sequence in sequences
+            )
+
+        padded = []
+        masks = []
+
+        for sequence in sequences:
+
+            sequence = sequence[:max_length]
+
+            padding_length = (
+                max_length - len(sequence)
+            )
+
+            padded_sequence = (
+                sequence
+                + [self.pad_token_id] * padding_length
+            )
+
+            mask = (
+                [1] * len(sequence)
+                + [0] * padding_length
+            )
+
+            padded.append(padded_sequence)
+            masks.append(mask)
+
+        return padded, masks

src/task/config.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+
+@dataclass
+class ModelConfig:
+
+    context_length: int = 64
+
+    d_model: int = 64
+    d_ff: int = 256
+    num_layers: int = 2
+
+    dropout: float = 0.1
+
+
+@dataclass
+class TrainingConfig:
+
+    batch_size: int = 32
+    learning_rate: float = 0.001
+    num_steps: int = 1000
+    eval_interval: int = 100
+    eval_batches: int = 20

src/task/inference.py

@@ -0,0 +1,43 @@
+import torch
+
+from src.infra.generation import generator
+
+
+class TextGenerator:
+
+    def __init__(
+        self,
+        model,
+        tokenizer,
+        context_length,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.context_length = context_length
+
+        self.model.eval()
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompt,
+        max_new_tokens=300,
+        temperature=0.7,
+    ):
+
+        prompt_tokens = torch.tensor(
+            [self.tokenizer.encode(prompt)],
+            dtype=torch.long,
+        )
+
+        generated_tokens = generator(
+            tokens=prompt_tokens,
+            generator=self.model,
+            context_length=self.context_length,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+        )
+
+        return self.tokenizer.decode(
+            generated_tokens[0].tolist()
+        )

src/task/model.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+
+from src.transformer.transformer import Transformer
+from src.transformer.input_representation import InputRepresentation
+
+class TinyTransformerLM(nn.Module):
+
+    def __init__(
+        self,
+        config,
+        vocab_size,
+    ):
+        super().__init__()
+
+        self.embedding = nn.Embedding(
+            vocab_size,
+            config.d_model,
+        )
+
+        self.input_representation = InputRepresentation(
+            vocab_size=vocab_size,
+            context_length=config.context_length,
+            d_model=config.d_model,
+            )
+
+        self.transformer = Transformer(
+            d_model=config.d_model,
+            d_ff=config.d_ff,
+            context_length=config.context_length,
+            num_layers=config.num_layers,
+            dropout=config.dropout,
+        )
+
+        self.head = nn.Linear(
+            config.d_model,
+            vocab_size,
+        )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+    ) -> torch.Tensor:
+
+        x = self.input_representation(tokens)
+
+        x = self.transformer(x)
+
+        logits = self.head(x)
+
+        return logits
+
+
+def load_model(
+    checkpoint_path,
+    config,
+    vocab_size,
+):
+    model = TinyTransformerLM(
+        config=config,
+        vocab_size=vocab_size,
+    )
+
+    checkpoint = torch.load(
+        checkpoint_path,
+        map_location="cpu",
+    )
+
+    model.load_state_dict(checkpoint)
+
+    model.eval()
+
+    return model

src/task/training.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn as nn
+
+from src.infra.dataset import (AutoRegressiveDataset, get_batch)
+
+
+@torch.no_grad()
+def evaluate(
+    model,
+    dataset: AutoRegressiveDataset,
+    model_config,
+    training_config,
+):
+    model.eval()
+
+    loss_function = nn.CrossEntropyLoss()
+
+    losses = []
+
+    for _ in range(training_config.eval_batches):
+
+        x, y = get_batch(
+            dataset=dataset,
+            batch_size=training_config.batch_size,
+        )
+
+        logits = model(x)
+
+        loss = loss_function(
+            logits.view(-1, logits.size(-1)),
+            y.view(-1),
+        )
+
+        losses.append(loss.item())
+
+    model.train()
+
+    return sum(losses) / len(losses)
+
+
+
+
+
+
+def train(
+    model,
+    train_dataset: AutoRegressiveDataset,
+    val_dataset: AutoRegressiveDataset,
+    model_config,
+    training_config,
+):
+    loss_function = nn.CrossEntropyLoss()
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=training_config.learning_rate,
+    )
+
+    model.train()
+
+    for step in range(training_config.num_steps):
+
+        if step % training_config.eval_interval == 0:
+
+            train_loss = evaluate(
+                model=model,
+                dataset=train_dataset,
+                model_config=model_config,
+                training_config=training_config,
+            )
+
+            val_loss = evaluate(
+                model=model,
+                dataset=val_dataset,
+                model_config=model_config,
+                training_config=training_config,
+            )
+
+            print(
+                f"step {step:4d} | "
+                f"train {train_loss:.4f} | "
+                f"val {val_loss:.4f}"
+            )
+
+        x, y = get_batch(
+            dataset=train_dataset,
+            batch_size=training_config.batch_size,
+        )
+
+        logits = model(x)
+
+        loss = loss_function(
+            logits.view(-1, logits.size(-1)),
+            y.view(-1),
+        )
+
+        optimizer.zero_grad()
+
+        loss.backward()
+
+        optimizer.step()
+
+    return model

src/train.py

@@ -0,0 +1,172 @@
+from pathlib import Path
+
+import torch
+
+from src.infra.tokenizer import CharacterTokenizer
+from src.infra.dataset import AutoRegressiveDataset, split_tokens
+
+from src.task.config import ModelConfig,TrainingConfig
+
+from src.task.model import TinyTransformerLM
+from src.task.training import train, evaluate
+
+
+
+def separator(title: str) -> None:
+    print("\n" + "=" * 60)
+    print(title)
+    print("=" * 60)
+
+
+# --------------------------------------------------
+# Configuration
+# --------------------------------------------------
+
+separator("Configuration")
+
+torch.manual_seed(42)
+
+model_config = ModelConfig()
+training_config = TrainingConfig()
+
+print(model_config)
+print(training_config)
+
+
+# --------------------------------------------------
+# Load corpus
+# --------------------------------------------------
+
+separator("Corpus")
+
+data_path = Path("data/tiny_shakespear.txt")
+
+text = data_path.read_text(
+    encoding="utf-8",
+)
+
+print(f"Corpus length: {len(text):,} characters")
+
+
+# --------------------------------------------------
+# Tokenization
+# --------------------------------------------------
+
+separator("Tokenization")
+
+tokenizer = CharacterTokenizer(text)
+
+tokens = torch.tensor(
+    tokenizer.encode(text),
+    dtype=torch.long,
+)
+
+print(f"Vocabulary size: {tokenizer.vocab_size}")
+print(f"Token count: {len(tokens):,}")
+
+
+# --------------------------------------------------
+# Train / validation / test split
+# --------------------------------------------------
+
+separator("Data Split")
+
+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):,}")
+
+
+# --------------------------------------------------
+# Datasets
+# --------------------------------------------------
+
+separator("Datasets")
+
+train_dataset = AutoRegressiveDataset(
+    tokens=train_tokens,
+    context_length=model_config.context_length,
+)
+
+val_dataset = AutoRegressiveDataset(
+    tokens=val_tokens,
+    context_length=model_config.context_length,
+)
+
+test_dataset = AutoRegressiveDataset(
+    tokens=test_tokens,
+    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):,}")
+
+
+# --------------------------------------------------
+# Model
+# --------------------------------------------------
+
+separator("Model")
+
+model = TinyTransformerLM(
+    config=model_config,
+    vocab_size=tokenizer.vocab_size,
+)
+
+print(model)
+
+
+# --------------------------------------------------
+# Training
+# --------------------------------------------------
+
+separator("Training")
+
+train(
+    model=model,
+    train_dataset=train_dataset,
+    val_dataset=val_dataset,
+    model_config=model_config,
+    training_config=training_config,
+)
+
+
+# --------------------------------------------------
+# Final test evaluation
+# --------------------------------------------------
+
+separator("Test Evaluation")
+
+test_loss = evaluate(
+    model=model,
+    dataset=test_dataset,
+    model_config=model_config,
+    training_config=training_config,
+)
+
+print(f"Test loss: {test_loss:.4f}")
+
+
+# --------------------------------------------------
+# Save checkpoint
+# --------------------------------------------------
+
+separator("Saving")
+
+checkpoint_dir = Path("checkpoints")
+checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+checkpoint_path = (
+    checkpoint_dir / "tiny-sentence-transformer.pt"
+)
+
+torch.save(
+    model.state_dict(),
+    checkpoint_path,
+)
+
+print(f"Checkpoint saved to: {checkpoint_path}")

src/transformer/attention.py

@@ -0,0 +1,277 @@
+"""
+Single-head causal self-attention.
+
+Learns relationships between tokens in the input sequence.
+
+Input Representation (X)
+        │
+        ▼
+    Q = X @ Wq
+    K = X @ Wk
+    V = X @ Wv
+        │
+        ▼
+    Scores = Q @ Kᵀ
+        │
+        ▼
+    Scale by √d
+        │
+        ▼
+    Apply Causal Mask
+        │
+        ▼
+    Softmax
+        │
+        ▼
+    Attention Weights
+        │
+        ▼
+Output = Attention Weights @ V
+
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class Attention(nn.Module):
+    """
+    Single-head causal self-attention.
+
+    Input:
+        (batch_size, context_length, d_model)
+
+    Output:
+        (batch_size, context_length, d_model)
+    """
+
+    def __init__(self, d_model: int, context_length: int,):
+        super().__init__()
+
+        # --------------------------------------------------
+        # Learnable projection layers.
+        #
+        # Each nn.Linear layer contains a learnable weight
+        # matrix that projects the input representation into
+        # a different vector space.
+        #
+        # Transformer paper notation:
+        #
+        #   Query Projection -> Wq
+        #   Key Projection   -> Wk
+        #   Value Projection -> Wv
+        #
+        # During the forward pass:
+        #
+        #   Q = X @ Wq
+        #   K = X @ Wk
+        #   V = X @ Wv
+        #
+        # In PyTorch, the multiplication is performed by
+        # calling the corresponding nn.Linear layer.
+        # --------------------------------------------------
+
+        self.query = nn.Linear(
+            in_features=d_model,
+            out_features=d_model,
+            bias=False,
+        )
+
+        self.key = nn.Linear(
+            in_features=d_model,
+            out_features=d_model,
+            bias=False,
+        )
+
+        self.value = nn.Linear(
+            in_features=d_model,
+            out_features=d_model,
+            bias=False,
+        )
+
+        # --------------------------------------------------
+        # Causal Mask
+        #
+        # Prevents tokens from attending to future tokens.
+        #
+        # The mask is created once during initialization
+        # and reused during every forward pass.
+        # --------------------------------------------------
+
+        self.register_buffer(
+            "mask",
+            torch.triu(
+                torch.ones(
+                    context_length,
+                    context_length,
+                    dtype=torch.bool,
+                ),
+                diagonal=1,
+            ),
+        )
+
+
+
+    def forward(self, 
+            x: torch.Tensor, 
+            padding_mask: torch.Tensor | None = None,
+            ) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        x
+            Input representation.
+
+            Shape:
+                (batch_size, context_length, d_model)
+
+        padding_mask
+            Optional mask indicating which positions contain
+            real tokens.
+
+            True  -> real token
+            False -> padding
+
+            Shape:
+                (batch_size, context_length)
+        """
+
+        # --------------------------------------------------
+        # Step 1
+        # Create Query (Q), Key (K) and Value (V)
+        #
+        # Mathematically:
+        #
+        #   Q = X @ Wq
+        #   K = X @ Wk
+        #   V = X @ Wv
+        #
+        # where X is the input representation and Wq, Wk and
+        # Wv are the learnable weight matrices contained in
+        # the projection layers above.
+        # --------------------------------------------------
+
+        Q = self.query(x)
+        K = self.key(x)
+        V = self.value(x)
+
+        # --------------------------------------------------
+        # Step 2
+        # Compute Attention Scores
+        #
+        # Mathematically:
+        #
+        #   Scores = Q @ Kᵀ
+        #
+        # Every Query vector is compared with every Key vector
+        # using the dot product, producing one attention score
+        # for every Query–Key pair.
+        #
+        # scores[i, j] answers:
+        #
+        #   "How much should token i attend to token j?"
+        # --------------------------------------------------
+
+        scores = Q @ K.transpose(-2, -1)
+
+        # --------------------------------------------------
+        # Step 3
+        # Scale Attention Scores
+        #
+        # Mathematically:
+        #
+        #           Q @ Kᵀ
+        # Scores = --------
+        #            √d
+        #
+        # where d is the embedding dimension (d_model).
+        #
+        # Scaling keeps the magnitude of the attention
+        # scores relatively stable regardless of the
+        # embedding dimension.
+        # --------------------------------------------------
+
+        scores = scores / math.sqrt(Q.size(-1))
+
+        # --------------------------------------------------
+        # Step 4
+        # Apply Causal Mask
+        #
+        # Prevent each token from attending to future tokens.
+        #
+        # The masked positions are assigned negative infinity.
+        # After applying softmax, they will receive a
+        # probability of zero.
+        # --------------------------------------------------
+
+        context_length = scores.size(-1)
+
+        # --------------------------------------------------
+        # Step 4
+        # Apply Causal Mask
+        # "Don't look into the future."
+        # --------------------------------------------------
+
+        scores = scores.masked_fill(
+            self.mask[:context_length, :context_length],
+            float("-inf"),
+            )
+
+        # --------------------------------------------------
+        # Apply Padding Mask
+        #
+        # True  -> real token
+        # False -> PAD
+        #
+        # We mask the KEY positions, meaning:
+        #
+        #   "Do not allow any token to attend
+        #    to a PAD token."
+        # --------------------------------------------------
+
+        if padding_mask is not None:
+            scores = scores.masked_fill(
+                ~padding_mask[:, None, :],
+                float("-inf"),
+            )
+        # --------------------------------------------------
+        # Step 5
+        # Apply Softmax
+        #
+        # Convert the attention scores into probabilities.
+        #
+        # Softmax is applied independently to each row of
+        # the attention score matrix so that every row sums
+        # to one.
+        #
+        # Each probability represents how much attention a
+        # Query token pays to each Key token.
+        # --------------------------------------------------
+
+        attention = torch.softmax(
+            scores,
+            dim=-1,
+        )
+
+        # --------------------------------------------------
+        # Step 6
+        # Compute the Weighted Sum
+        #
+        # Mathematically:
+        #
+        #   Output = Attention Weights @ V
+        #
+        # The attention weights determine how much
+        # information to retrieve from every Value vector.
+        #
+        # Each output vector is a weighted combination of
+        # all Value vectors visible to the current Query.
+        # --------------------------------------------------
+
+        output = attention @ V
+
+        return output

src/transformer/embedding.py

@@ -0,0 +1,31 @@
+"""
+Token embedding layer.
+
+Transforms token IDs into dense, learnable vectors.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class Embedding(nn.Module):
+    """
+    Token Embedding
+
+    Input:
+        (batch_size, context_length)
+
+    Output:
+        (batch_size, context_length, d_model)
+    """
+
+    def __init__(self, vocab_size: int, d_model: int):
+        super().__init__()
+
+        self.embedding = nn.Embedding(
+            num_embeddings=vocab_size,
+            embedding_dim=d_model
+        )
+
+    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.embedding(token_ids)

src/transformer/feed_forward.py

@@ -0,0 +1,26 @@
+"""
+Attention
+    → tokens exchange information
+
+Feed Forward
+    → each token independently transforms
+      its context-aware representation
+"""
+
+import torch
+import torch.nn as nn
+
+
+class FeedForward(nn.Module):
+
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Linear(d_model, d_ff),       #d_ff is the hidden dimension of the feed forward network, usually larger than d_model
+            nn.GELU(),                      #Activation function.. Gaussian Error Linear Unit (GELU) is a smooth approximation of the ReLU function, which allows for better gradient flow during training.
+            nn.Linear(d_ff, d_model),       #collapses the hidden dimension back to the original model dimension, ensuring that the output has the same shape as the input
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)

src/transformer/input_representation.py

@@ -0,0 +1,66 @@
+"""
+Input representation layer.
+
+Combines token embeddings and positional embeddings to produce
+the representation that enters the transformer.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from .embedding import Embedding
+from .positional_embedding import PositionalEmbedding
+
+
+class InputRepresentation(nn.Module):
+    """
+    Input Representation
+
+    Input:
+        token_ids
+            Shape:
+                (batch_size, context_length)
+
+    Output:
+        input_representation
+            Shape:
+                (batch_size, context_length, d_model)
+
+    Formula:
+        Input Representation =
+            Token Embedding +
+            Positional Embedding
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        context_length: int,
+        d_model: int,
+    ) -> None:
+        super().__init__()
+
+        self.token_embedding = Embedding(
+            vocab_size=vocab_size,
+            d_model=d_model,
+        )
+
+        self.positional_embedding = PositionalEmbedding(
+            context_length=context_length,
+            d_model=d_model,
+        )
+
+    def forward(
+        self,
+        token_ids: torch.Tensor,
+    ) -> torch.Tensor:
+
+        token_vectors = self.token_embedding(token_ids)
+
+        position_vectors = self.positional_embedding(token_ids)
+
+        return token_vectors + position_vectors
+
+

src/transformer/positional_embedding.py

@@ -0,0 +1,41 @@
+"""
+Positional embedding layer.
+
+Converts token positions into learnable embedding vectors.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class PositionalEmbedding(nn.Module):
+    """
+    Learned positional embeddings.
+
+    Input:
+        token_ids: (batch_size, context_length)
+
+    Output:
+        (batch_size, context_length, d_model)
+    """
+
+    def __init__(self, context_length: int, d_model: int):
+        super().__init__()
+
+        self.embedding = nn.Embedding(
+            num_embeddings=context_length,
+            embedding_dim=d_model
+        )
+
+    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
+
+        batch_size, context_length = token_ids.shape
+
+        positions = torch.arange(
+            context_length,
+            device=token_ids.device
+        )
+
+        positions = positions.unsqueeze(0).expand(batch_size, context_length)
+
+        return self.embedding(positions)

src/transformer/transformer.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn as nn
+
+from src.transformer.transformer_block import TransformerBlock
+
+
+class Transformer(nn.Module):
+
+    def __init__(
+        self,
+        d_model: int,
+        d_ff: int,
+        context_length: int,
+        num_layers: int,
+        dropout = 0.1,
+    ):
+        super().__init__()
+
+        self.blocks = nn.ModuleList([
+            TransformerBlock(
+                d_model=d_model,
+                d_ff=d_ff,
+                context_length=context_length,
+                dropout=dropout,
+            )
+            for _ in range(num_layers)
+        ])
+
+    def forward(self, 
+                x: torch.Tensor, 
+                padding_mask: torch.Tensor | None = None,) -> torch.Tensor:
+
+        for block in self.blocks:
+            x = block(x,padding_mask=padding_mask,)
+
+        return x

src/transformer/transformer_block.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+
+from src.transformer.attention import Attention
+from src.transformer.feed_forward import FeedForward
+
+
+class TransformerBlock(nn.Module):
+
+    def __init__(self, d_model: int, d_ff: int, context_length: int, dropout=0.1,):
+        super().__init__()
+
+        self.attention = Attention(d_model, context_length)
+        self.feed_forward = FeedForward(d_model, d_ff)
+
+        self.layer_norm_1 = nn.LayerNorm(d_model)
+        self.layer_norm_2 = nn.LayerNorm(d_model)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, padding_mask: torch.Tensor | None = None,) -> torch.Tensor:
+
+        attention_output = self.attention(x, padding_mask=padding_mask)
+
+        x = self.layer_norm_1(
+            x + self.dropout(attention_output)
+        )
+
+        ffn_output = self.feed_forward(x)
+
+        x = self.layer_norm_2(
+            x + self.dropout(ffn_output)
+        )
+
+        return x