selfies-tokenizer 0.1.0

Super fast tokenizer for SELFIES molecular representations with ML-ready batch processing

SELFIES Tokenizer

A super fast tokenizer for SELFIES (SELF-referencIng Embedded Strings) molecular representations with ML-ready batch processing and vocabulary management.

Features

• Blazing fast: Optimized regex-based tokenization (~109k molecules/sec on real datasets)
• ML-ready: Batch processing, vocabulary management, and encoding to indices
• Dataset packaging: Save/load encoded datasets with full metadata (vocab, data, index mapping, statistics)
• Smart max_len selection: Interactive suggest_len() suggests optimal sequence lengths based on coverage (100%, 90%, 75%)
• Special tokens: Automatic <start> and <end> token insertion for sequence models
• Smart decoding: Stops at <end> token and removes all special tokens automatically
• Padding & Truncation: Built-in max_len support with <pad> always at index 0
• Metadata management: Saves max_len and vocab_len in vocabulary file
• Progress bars: Built-in tqdm support for large dataset processing
• Flexible encoding: Return token strings or indices with return_str flag
• Vocabulary management: Build, save, and load vocabularies for consistent encoding
• Multiple methods: Regex (fastest), manual parsing, or selfies library
• Simple API: Easy to use with sensible defaults
• Well-tested: Tested on 768k real SELFIES molecules
• Zero dependencies: Core functionality requires no external libraries (numpy for dataset saving, tqdm optional)

Installation

pip install selfies-tokenizer  # Coming soon to PyPI

For now, install from source:

git clone <repo>
cd selfies-tokenizer
pip install -e .

Quick Start

Basic Tokenization

from selfies_tokenizer import tokenize

# Simple usage
tokens = tokenize('[C][=O]')
# Returns: ['[C]', '[=O]']

ML Workflow with Special Tokens

from selfies_tokenizer import SELFIESTokenizer

# 1. Build vocabulary from training data with default max_len
tokenizer = SELFIESTokenizer(
    max_len=15,  # Saved in metadata
    vocab_path='./vocab.json',
    refresh_dict=True
)
train_data = ['[C][=O][OH]', '[N][C][C]', '[C][Branch1][C][O]']
tokenizer.fit(train_data, show_progress=True)  # Progress bar for large datasets
tokenizer.save_vocab()  # Saves max_len and vocab_len in metadata

# Note: <pad> is always at index 0
print(tokenizer.token2idx['<pad>'])  # 0

# 2. Encode to indices (adds <start> and <end>, then pads)
indices = tokenizer.encode('[C][=O]')
# Returns: [2, 6, 5, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
# Format: <start>, [C], [=O], <end>, <pad>... (length 15)

# 3. Batch processing with uniform length
batch_indices = tokenizer.encode(['[C][=O]', '[N]', '[C][OH]'], show_progress=True)
# All sequences have same length (15 from metadata)

# 4. Load existing vocabulary for inference (max_len auto-loaded)
inference_tokenizer = SELFIESTokenizer(vocab_path='./vocab.json')
test_indices = inference_tokenizer.encode('[C][=O]')  # Uses max_len=15 from metadata

# 5. Decode predictions (stops at <end>, removes special tokens)
decoded = inference_tokenizer.decode([2, 6, 5, 1, 0, 0, 0])
# Returns: '[C][=O]' (clean output, special tokens removed!)

# 6. Batch decode
decoded_batch = inference_tokenizer.decode(batch_indices)
# Returns: ['[C][=O]', '[N]', '[C][OH]'] (all clean)

Smart max_len Selection with suggest_len()

Let the tokenizer analyze your data and suggest optimal max_len values based on coverage:

from selfies_tokenizer import SELFIESTokenizer

# Build vocabulary
tokenizer = SELFIESTokenizer(vocab_path='./vocab.json', refresh_dict=True)
tokenizer.fit(train_data)

# Get smart max_len suggestions (interactive)
max_len = tokenizer.suggest_len(train_data, show_progress=True)

# Output:
# ======================================================================
# SEQUENCE LENGTH ANALYSIS
# ======================================================================
#
# Dataset: 1000 sequences
# Min length: 10
# Median length: 37
# Max length: 73
#
# ----------------------------------------------------------------------
# SUGGESTED MAX_LEN OPTIONS
# ----------------------------------------------------------------------
#
# [1] 100% coverage (max_len=73)
#     → 0 sequences truncated (0.0%)
#     → All sequences preserved completely
#
# [2] 90% coverage (max_len=55)
#     → 93 sequences truncated (9.3%)
#     → 907 sequences preserved (90.7%)
#
# [3] 75% coverage (max_len=48)
#     → 249 sequences truncated (24.9%)
#     → 751 sequences preserved (75.1%)
#
# [4] Custom max_len
#
# ----------------------------------------------------------------------
# Select option [1-4]: 2
# ✓ Selected: max_len=55 (90% coverage)

# Now encode with the selected max_len
encoded_data = tokenizer.encode(train_data, max_len=max_len)

Benefits:

• Data-driven: Analyzes your actual sequence lengths
• Trade-off visualization: See exactly how many sequences get truncated
• Interactive: Choose the coverage that fits your needs
• Custom option: Enter your own max_len with automatic coverage calculation

Save Dataset to NumPy with Metadata

Save your encoded dataset with full metadata for ML training:

from selfies_tokenizer import SELFIESTokenizer
from selfies_tokenizer import save_encoded_dataset, load_encoded_dataset

# Build tokenizer
tokenizer = SELFIESTokenizer()
tokenizer.fit(train_data)

# Save complete dataset package
paths = save_encoded_dataset(
    selfies_data=train_data,
    tokenizer=tokenizer,
    save_dir='./datasets/my_dataset',
    max_len=50,
    show_progress=True
)

# Creates directory structure:
# ./datasets/my_dataset/
#     ├── vocab.json              # Tokenizer vocabulary
#     ├── encoded_data.npy        # Encoded sequences (shape: [N, 50])
#     ├── index_mapping.json      # Maps array index -> SELFIES string
#     └── metadata.json           # Timestamp, settings, statistics

# Later, load everything back
dataset = load_encoded_dataset('./datasets/my_dataset')

# Ready for ML!
X = dataset['encoded_data']        # numpy array
tokenizer = dataset['tokenizer']   # SELFIESTokenizer instance
metadata = dataset['metadata']     # dict with creation time, settings, etc.

# Use with PyTorch/TensorFlow
import torch
X_tensor = torch.tensor(X, dtype=torch.long)

What gets saved:

• vocab.json: Complete tokenizer vocabulary with all settings
• encoded_data.npy: Encoded sequences as numpy array (int32)
• index_mapping.json: Maps each array index to its SELFIES string
• metadata.json: Creation timestamp, encoding settings, dataset statistics, coverage info

Metadata includes:

• Creation timestamp
• Dataset size and statistics (min/max/mean/median sequence lengths)
• Encoding settings (max_len, padding, truncation, special tokens)
• Tokenizer configuration (vocab size, method)
• Coverage percentage (how many sequences were not truncated)
• Array shape and dtype

API Reference

SELFIESTokenizer

Initialization

tokenizer = SELFIESTokenizer(
    method='auto',              # Tokenization method: 'auto', 'regex', 'manual', 'selfies_lib'
    vocab_path=None,            # Path to vocabulary JSON file
    refresh_dict=False,         # Force rebuild vocab even if exists
    special_tokens=None,        # Custom special tokens (default: ['<pad>', '<unk>', '<start>', '<end>'])
    max_len=None                # Default max sequence length (saved in metadata)
)

Parameters:

• method: Tokenization algorithm

  • 'auto': Automatically select fastest (default: regex)
  • 'regex': Pre-compiled regex (fastest)
  • 'manual': Manual string parsing
  • 'selfies_lib': Use selfies library (requires: pip install selfies)

• vocab_path: Path to vocabulary file

  • If file exists and refresh_dict=False: Loads existing vocabulary (including max_len)
  • If file doesn't exist: Will save vocabulary here when calling save_vocab()

• refresh_dict: Force vocabulary rebuild

  • True: Ignore existing vocabulary file, start fresh
  • False: Load existing vocabulary if available

• special_tokens: List of special tokens

  • Default: ['<pad>', '<unk>', '<start>', '<end>']
  • Customize for your use case: ['<pad>', '<mask>', '<cls>', '<sep>']
  • <pad> is always placed at index 0

• max_len: Default maximum sequence length

  • Saved in metadata when calling save_vocab()
  • Auto-loaded when loading vocabulary
  • Used as default in encode() if not specified

Methods

fit(selfies_data, show_progress=False)

Build vocabulary from SELFIES data.

tokenizer.fit(['[C][=O]', '[N][C]', '[C][OH]'])
# Builds vocabulary from training data

# With progress bar for large datasets
tokenizer.fit(large_dataset, show_progress=True)

Parameters:

• selfies_data: str or List[str] - SELFIES string(s) to build vocabulary from
• show_progress: bool - Show progress bar (requires tqdm)

Returns: self (for method chaining)

---

suggest_len(selfies_data, add_special_tokens=True, show_progress=False)

Analyze sequence lengths and suggest optimal max_len values with interactive coverage selection.

# Analyze data and get suggestions
max_len = tokenizer.suggest_len(train_data, show_progress=True)

# Interactive prompt shows:
# - Min/median/max lengths
# - 100% coverage (no truncation)
# - 90% coverage
# - 75% coverage
# - Custom max_len option

# Returns user's chosen max_len

Parameters:

• selfies_data: str or List[str] - SELFIES string(s) to analyze
• add_special_tokens: bool (default: True) - Account for <start> and <end> in length calculations
• show_progress: bool (default: False) - Show progress bar during analysis

Returns: int - User's selected max_len value (or None if cancelled)

Interactive Options:

1. 100% coverage: max_len = longest sequence (no truncation)
2. 90% coverage: max_len at 90th percentile (~10% truncated)
3. 75% coverage: max_len at 75th percentile (~25% truncated)
4. Custom: Enter your own max_len with automatic coverage calculation

Use Case: Perfect for finding the optimal trade-off between sequence length and truncation. Longer max_len preserves more data but increases memory usage and computation. This tool helps you make an informed decision based on your actual data distribution.

---

encode(selfies_data, return_str=False, max_len=None, padding=True, truncation=True, show_progress=False, add_special_tokens=True)

Encode SELFIES to indices or token strings. Automatically adds <start> and <end> tokens. Supports both single and batch inputs with padding and truncation.

# Single string to indices (adds <start> and <end>)
indices = tokenizer.encode('[C][=O]', max_len=10)
# Returns: [2, 5, 4, 1, 0, 0, 0, 0, 0, 0]
# Format: <start>, [C], [=O], <end>, <pad>...

# Batch encoding with progress bar
batch_indices = tokenizer.encode(large_dataset, max_len=10, show_progress=True)

# Without special tokens
indices = tokenizer.encode('[C][=O]', max_len=10, add_special_tokens=False)
# Returns: [5, 4, 0, 0, 0, 0, 0, 0, 0, 0]

# Return string tokens instead of indices
tokens = tokenizer.encode('[C]', max_len=5, return_str=True)
# Returns: ['<start>', '[C]', '<end>', '<pad>', '<pad>']

Parameters:

• selfies_data: str or List[str] - SELFIES string(s) to encode
• return_str: bool (default: False) - Return token strings instead of indices
• max_len: int or None (default: None) - Maximum sequence length. Uses self.max_len if not specified
• padding: bool (default: True) - Pad sequences shorter than max_len with <pad> (index 0)
• truncation: bool (default: True) - Truncate sequences longer than max_len
• show_progress: bool (default: False) - Show progress bar (requires tqdm)
• add_special_tokens: bool (default: True) - Add <start> and <end> tokens

Returns:

• Single input + return_str=False: List[int] (token indices)
• Single input + return_str=True: List[str] (token strings)
• Batch input + return_str=False: List[List[int]] (batch of indices)
• Batch input + return_str=True: List[List[str]] (batch of tokens)

Notes:

• <pad> token is always at index 0 for efficient masking in ML models
• <start> and <end> tokens are automatically added by default
• Sequence format: [<start>, token1, token2, ..., <end>, <pad>, ...]
• When max_len is specified, all output sequences will have the same length
• Essential for batching in PyTorch/TensorFlow

---

decode(indices, skip_special_tokens=True)

Decode token indices back to SELFIES strings. Automatically stops at <end> token and removes special tokens.

# Single decode (stops at <end>, removes special tokens)
selfies = tokenizer.decode([2, 5, 4, 1, 0, 0])
# Input: <start>, [C], [=O], <end>, <pad>, <pad>
# Returns: '[C][=O]' (clean output!)

# Batch decode
batch_selfies = tokenizer.decode([[2, 5, 4, 1], [2, 6, 5, 1]])
# Returns: ['[C][=O]', '[N][C]']

# Keep special tokens if needed
selfies = tokenizer.decode([2, 5, 4, 1], skip_special_tokens=False)
# Returns: '<start>[C][=O]' (stops at <end>, but keeps other special tokens)

Parameters:

• indices: List[int] or List[List[int]] - Token indices to decode
• skip_special_tokens: bool (default: True) - Remove all special tokens from output

Returns:

• Single input: str (SELFIES string)
• Batch input: List[str] (batch of SELFIES strings)

Behavior:

• Always stops at <end> token (doesn't decode beyond it)
• Removes all special tokens by default (<start>, <end>, <pad>, <unk>)
• Perfect for decoding model predictions

---

save_vocab(path=None)

Save vocabulary to JSON file.

tokenizer.save_vocab('./vocab.json')
# Or use vocab_path from __init__
tokenizer.save_vocab()

Parameters:

• path: str (optional) - Path to save vocabulary. Uses vocab_path from init if not specified.

---

load_vocab(path=None)

Load vocabulary from JSON file.

tokenizer.load_vocab('./vocab.json')
# Or use vocab_path from __init__
tokenizer.load_vocab()

Parameters:

• path: str (optional) - Path to vocabulary file. Uses vocab_path from init if not specified.

Returns: self (for method chaining)

---

tokenize(selfies_string)

Low-level tokenization to strings only (no vocabulary needed).

tokens = tokenizer.tokenize('[C][=O]')
# Returns: ['[C]', '[=O]']

Parameters:

• selfies_string: str - SELFIES string to tokenize

Returns: List[str] - Token strings

---

Properties

vocab_size

Get the size of the vocabulary.

size = tokenizer.vocab_size
# Returns: int

---

Convenience Functions

tokenize(selfies_string, method='auto')

Quick tokenization without creating a tokenizer instance.

from selfies_tokenizer import tokenize

tokens = tokenize('[C][=O]')
# Returns: ['[C]', '[=O]']

Utility Functions

save_encoded_dataset(...)

Save encoded SELFIES dataset with complete metadata package.

from selfies_tokenizer import save_encoded_dataset

paths = save_encoded_dataset(
    selfies_data=train_data,        # SELFIES strings to encode
    tokenizer=tokenizer,             # Fitted tokenizer
    save_dir='./datasets/my_data',   # Output directory
    max_len=50,                      # Sequence length
    padding=True,                    # Pad sequences
    truncation=True,                 # Truncate long sequences
    add_special_tokens=True,         # Add <start> and <end>
    show_progress=True               # Show progress bar
)

Parameters:

• selfies_data: str or List[str] - SELFIES strings to encode and save
• tokenizer: SELFIESTokenizer - Fitted tokenizer instance
• save_dir: str - Directory to save dataset (created if doesn't exist)
• max_len: int or None - Maximum sequence length
• padding: bool - Pad sequences to max_len
• truncation: bool - Truncate sequences longer than max_len
• add_special_tokens: bool - Add <start> and <end> tokens
• show_progress: bool - Show progress bar
• vocab_filename: str - Vocabulary filename (default: 'vocab.json')
• data_filename: str - Data filename (default: 'encoded_data.npy')
• index_filename: str - Index mapping filename (default: 'index_mapping.json')
• metadata_filename: str - Metadata filename (default: 'metadata.json')

Returns: dict with paths to all saved files

Creates:

• vocab.json - Tokenizer vocabulary
• encoded_data.npy - Encoded sequences as numpy array (int32, shape: [N, max_len])
• index_mapping.json - Maps each array index to its SELFIES string
• metadata.json - Timestamp, settings, statistics, coverage info

---

load_encoded_dataset(...)

Load encoded SELFIES dataset with all metadata.

from selfies_tokenizer import load_encoded_dataset

dataset = load_encoded_dataset(
    load_dir='./datasets/my_data',
    load_tokenizer=True
)

# Access components
X = dataset['encoded_data']        # numpy array
tokenizer = dataset['tokenizer']   # SELFIESTokenizer
metadata = dataset['metadata']     # dict with all info
index_mapping = dataset['index_mapping']  # index -> SELFIES

Parameters:

• load_dir: str - Directory containing the dataset
• vocab_filename: str - Vocabulary filename (default: 'vocab.json')
• data_filename: str - Data filename (default: 'encoded_data.npy')
• index_filename: str - Index mapping filename (default: 'index_mapping.json')
• metadata_filename: str - Metadata filename (default: 'metadata.json')
• load_tokenizer: bool - If True, load and return tokenizer instance

Returns: dict containing:

• encoded_data: np.ndarray - Encoded sequences
• index_mapping: dict - Maps array index to SELFIES string
• metadata: dict - Full metadata (timestamp, settings, statistics)
• tokenizer: SELFIESTokenizer or None - Tokenizer instance (if load_tokenizer=True)

Padding and Truncation

The tokenizer provides built-in padding and truncation for ML training. The <pad> token is always at index 0, making it easy to create attention masks and work with PyTorch/TensorFlow.

Why Padding and Truncation?

Neural networks require fixed-length input sequences for batching. Variable-length SELFIES strings need to be:

• Padded: Short sequences extended to max_len with <pad> tokens
• Truncated: Long sequences cut to max_len

Basic Usage

from selfies_tokenizer import SELFIESTokenizer

tokenizer = SELFIESTokenizer()
tokenizer.fit(['[C][=O]', '[N][C][C]'])

# Pad short sequence
result = tokenizer.encode('[C]', max_len=4)
# Returns: [5, 0, 0, 0]  (padded with 0)

# Truncate long sequence
result = tokenizer.encode('[C][C][C][C][C]', max_len=3)
# Returns: [5, 5, 5]  (truncated)

# Batch with uniform length
batch = ['[C]', '[C][=O]', '[N][C][C]']
result = tokenizer.encode(batch, max_len=4)
# Returns: [[5, 0, 0, 0], [5, 4, 0, 0], [6, 5, 5, 0]]
# All sequences now have length 4!

Key Features

• <pad> always at index 0: Guaranteed for all vocabularies
• Batch processing: Entire batches padded/truncated in one call
• Flexible control: Enable/disable padding and truncation independently
• String support: Works with return_str=True for debugging

Control Flags

# Only padding, no truncation
tokenizer.encode('[C]', max_len=5, padding=True, truncation=False)

# Only truncation, no padding
tokenizer.encode('[C][C][C]', max_len=2, padding=False, truncation=True)

# Neither (variable length output)
tokenizer.encode('[C][=O]', max_len=5, padding=False, truncation=False)

PyTorch Example

import torch

# Encode with fixed length
data = ['[C][=O]', '[N]', '[C][C][C]']
encoded = tokenizer.encode(data, max_len=4)

# Convert to tensor
tensor = torch.tensor(encoded)
print(tensor.shape)  # torch.Size([3, 4])

# Create attention mask (1 for real tokens, 0 for padding)
attention_mask = (tensor != 0).long()

Performance

Benchmarked on a standard machine with 10,000 iterations:

Method	Simple Molecule	Complex Molecule	Long Chain (500 tokens)
Regex	0.74µs	1.33µs	60.11µs
Manual	0.81µs	3.55µs	141.08µs
Selfies Lib	1.24µs	3.18µs	195.86µs

The regex method is 2-3x faster than alternatives.

Examples

Example 1: Training a Model

from selfies_tokenizer import SELFIESTokenizer

# Load your training data
train_selfies = [
    '[C][=O][OH]',
    '[N][C][C][C]',
    '[C][Branch1][C][O][C]',
    # ... more training data
]

# Create tokenizer and build vocabulary
tokenizer = SELFIESTokenizer(vocab_path='./model_vocab.json', refresh_dict=True)
tokenizer.fit(train_selfies)
tokenizer.save_vocab()

print(f"Vocabulary size: {tokenizer.vocab_size}")

# Encode training data to indices
train_indices = tokenizer.encode(train_selfies)

# Now train your model with train_indices...

Example 2: Inference with Saved Vocabulary

from selfies_tokenizer import SELFIESTokenizer

# Load pre-built vocabulary
tokenizer = SELFIESTokenizer(vocab_path='./model_vocab.json')

# Encode new molecule
molecule = '[C][=C][C][=O]'
indices = tokenizer.encode(molecule)

# Pass to your trained model...
prediction = model.predict(indices)

# Decode prediction back to SELFIES
predicted_selfies = tokenizer.decode(prediction)
print(f"Predicted molecule: {predicted_selfies}")

Example 3: Batch Processing

from selfies_tokenizer import SELFIESTokenizer

tokenizer = SELFIESTokenizer(vocab_path='./vocab.json')

# Process large batch efficiently
large_batch = ['[C][=O]', '[N][C]', ...] * 1000  # 1000s of molecules

# Single call for entire batch
batch_indices = tokenizer.encode(large_batch)

# Process batch...

Example 4: Custom Special Tokens

from selfies_tokenizer import SELFIESTokenizer

# Use custom special tokens for BERT-style models
tokenizer = SELFIESTokenizer(
    special_tokens=['<pad>', '<mask>', '<cls>', '<sep>'],
    vocab_path='./bert_vocab.json'
)

tokenizer.fit(train_data)

# <mask> token is now in vocabulary
mask_idx = tokenizer.token2idx['<mask>']

Example 5: Unknown Token Handling

from selfies_tokenizer import SELFIESTokenizer

# Train with limited vocabulary
tokenizer = SELFIESTokenizer()
tokenizer.fit(['[C][=O]', '[N][C]'])

# Encode molecule with unknown token
result = tokenizer.encode('[S][=O]')  # [S] not in vocabulary

# [S] is mapped to <unk>
unk_idx = tokenizer.token2idx['<unk>']
# result will contain unk_idx for [S]

Running Tests

# Run basic tests
python test_tokenizer.py

# Run ML functionality tests
python test_tokenizer_ml.py

# Run padding and truncation tests
python test_padding.py

# Run performance benchmarks
python benchmark.py

# Run examples
python example.py          # Basic usage
python example_ml.py       # ML workflow
python example_padding.py  # Padding and truncation

How It Works

SELFIES strings are formatted with tokens enclosed in square brackets:

• Input: [C][=O]
• Output: ['[C]', '[=O]']

The tokenizer:

1. Tokenizes using pre-compiled regex pattern \[[^\]]+\] (fastest method)
2. Builds vocabulary from training data with special tokens (<pad> always at index 0)
3. Encodes by mapping tokens to indices using the vocabulary
4. Pads/Truncates sequences to max_len if specified (for ML batching)
5. Handles unknown tokens by mapping them to <unk>
6. Decodes by mapping indices back to tokens and joining them

Requirements

• Python 3.6+
• Optional: selfies library (only for selfies_lib method)

pip install selfies  # Optional

Use Cases

• Machine Learning: Train models on molecular representations
• Drug Discovery: Process large chemical databases efficiently
• Chemical Informatics: Fast tokenization for analysis pipelines
• Model Deployment: Consistent vocabulary for training and inference

License

MIT

Contributing

Contributions welcome! Please feel free to submit a Pull Request.

Roadmap

• Add padding and truncation utilities
• Add PyTorch/TensorFlow dataset adapters
• Add CLI tool for vocabulary building
• Publish to PyPI
• Add sequence alignment utilities
• Add attention mask generation utilities