embedding = nn.Embedding(50257, 768)

input_ids = torch.tensor([[40, 952, 37167]])
vectors = embedding(input_ids)

#!/usr/bin/env python
"""Chapter 2: character tokenizer, a tiny BPE, and embeddings."""

from __future__ import annotations

import argparse
from collections import Counter
from pathlib import Path

import torch
from torch import nn


class CharacterTokenizer:
    def __init__(self) -> None:
        self.tokens: list[str] = []
        self.token_to_id: dict[str, int] = {}

    def fit(self, text: str) -> None:
        self.tokens = sorted(set(text))
        self.token_to_id = {token: index for index, token in enumerate(self.tokens)}

    def encode(self, text: str) -> list[int]:
        return [self.token_to_id[token] for token in text]

    def decode(self, ids: list[int]) -> str:
        return "".join(self.tokens[index] for index in ids)

    @property
    def vocab_size(self) -> int:
        return len(self.tokens)


class TinyBPE:
    """A deliberately small byte/character-level BPE implementation."""

    def __init__(self, num_merges: int = 50) -> None:
        self.num_merges = num_merges
        self.merges: list[tuple[str, str]] = []
        self.vocab: list[str] = []
        self.token_to_id: dict[str, int] = {}

    @staticmethod
    def _count_pairs(tokens: list[str]) -> Counter[tuple[str, str]]:
        return Counter(zip(tokens[:-1], tokens[1:]))

    @staticmethod
    def _merge(tokens: list[str], pair: tuple[str, str], replacement: str) -> list[str]:
        output: list[str] = []
        index = 0
        while index < len(tokens):
            if index + 1 < len(tokens) and (tokens[index], tokens[index + 1]) == pair:
                output.append(replacement)
                index += 2
            else:
                output.append(tokens[index])
                index += 1
        return output

    def fit(self, text: str) -> None:
        tokens = list(text)
        base_vocab = sorted(set(tokens))
        for _ in range(self.num_merges):
            pairs = self._count_pairs(tokens)
            if not pairs:
                break
            best_pair = max(pairs.items(), key=lambda item: (item[1], item[0]))[0]
            replacement = best_pair[0] + best_pair[1]
            tokens = self._merge(tokens, best_pair, replacement)
            self.merges.append(best_pair)
            if replacement not in base_vocab:
                base_vocab.append(replacement)
        self.vocab = base_vocab
        self.token_to_id = {token: index for index, token in enumerate(self.vocab)}

    def encode(self, text: str) -> list[str]:
        tokens = list(text)
        for pair in self.merges:
            replacement = pair[0] + pair[1]
            tokens = self._merge(tokens, pair, replacement)
        return tokens

    def encode_ids(self, text: str) -> list[int]:
        return [self.token_to_id[token] for token in self.encode(text)]

    @property
    def vocab_size(self) -> int:
        return len(self.vocab)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--text-file", type=Path, default=None)
    parser.add_argument("--merges", type=int, default=50)
    parser.add_argument("--embedding-dim", type=int, default=16)
    parser.add_argument("--seed", type=int, default=1337)
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    if args.text_file is not None:
        text = args.text_file.read_text(encoding="utf-8")[:50000]
    else:
        text = (
            "The transformer architecture uses self-attention. "
            "Self-attention lets each token attend to other tokens. "
            "Tokens become vectors, and vectors become context-aware representations. "
            "现代大语言模型把文本切成 token，再把 token 映射成向量。"
        )

    char_tokenizer = CharacterTokenizer()
    char_tokenizer.fit(text)

    bpe = TinyBPE(num_merges=args.merges)
    bpe.fit(text)

    char_ids = char_tokenizer.encode(text)
    bpe_ids = bpe.encode_ids(text)
    compression = len(char_ids) / max(1, len(bpe_ids))

    print(f"characters:              {len(text)}")
    print(f"character vocab size:    {char_tokenizer.vocab_size}")
    print(f"BPE vocab size:          {bpe.vocab_size}")
    print(f"character token count:   {len(char_ids)}")
    print(f"BPE token count:         {len(bpe_ids)}")
    print(f"compression ratio:       {compression:.3f}x")
    print(f"first 10 BPE merges:     {bpe.merges[:10]}")

    embedding = nn.Embedding(bpe.vocab_size, args.embedding_dim)
    sample = torch.tensor([bpe_ids[:12]], dtype=torch.long)
    vectors = embedding(sample)
    print(f"sample ids:              {sample.tolist()}")
    print(f"embedding output shape:  {tuple(vectors.shape)}")
    print(f"embedding table shape:   {tuple(embedding.weight.shape)}")


if __name__ == "__main__":
    main()