babybert 0.1.1

Minimal BERT implementation in PyTorch

---

Minimal implementation of the BERT architecture proposed by Devlin et al. using the PyTorch library. This implementation focuses on simplicity and readability, so the model code is not optimized for inference or training efficiency. BabyBERT can be fine-tuned for downstream tasks such as named-entity recognition (NER), sentiment classification, or question answering (QA).

See the roadmap below for my future plans for this library!

📦 Installation

pip install babybert

🚀 Quickstart

The following example demonstrates how to tokenize text, instantiate a BabyBERT model, and obtain contextual embeddings:

from babybert.tokenizer import WordPieceTokenizer
from babybert.model import BabyBERTConfig, BabyBERT

# Load a pretrained tokenizer and encode a text
tokenizer = WordPieceTokenizer.from_pretrained("toy-tokenizer")
encoded = tokenizer.batch_encode(["Hello, world!"])

# Initialize an untrained BabyBERT model
model_cfg = BabyBERTConfig.from_preset(
  "tiny", vocab_size=tokenizer.vocab_size, block_size=len(encoded['token_ids'][0])
)
model = BabyBERT(model_cfg)

# Obtain contextual embeddings
hidden = model(**encoded)
print(hidden)

[!TIP] For more usage examples, check out the examples/ directory!

🗺️ Roadmap

Model Implementation

• Build initial model implementation
• Write trainer class
• Create custom WordPiece tokenizer
• Introduce more parameter configurations
• Set up pretrained model checkpoints

Usage Examples

• Pretraining
• Sentiment classification
• Named entity recognition
• Question answering