Byte Latent Transformer

My pet implementation of Meta's Byte Latent Transformer (BLT) following the research paper: https://arxiv.org/pdf/2412.09871

Setup

python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

Training: train_blt.py

Train the full BLT model

python train_blt.py \
    --data_dir ./data \
    --output_dir ./models \
    --d_model 512 \
    --num_layers 6 \
    --local_layers 3 \
    --max_length 256 \
    --batch_size 8 \
    --blt_epochs 10 \
    --dropout 0.2

Output:

• models/blt_model_end_to_end.safetensors - Complete trained BLT model
• models/blt_config_end_to_end.json - Model configuration
• plots/training_curves.png - Training visualizations

Embeddings: generate_embeddings.py

Generate embeddings for files using a trained BLT model:

# Generate embeddings for all files
python generate_embeddings.py \
  --model_dir ./models \
  --data_dir ./documents \
  --output embeddings.safetensors

# With similarity matrix and custom patterns
python generate_embeddings.py \
  --model_dir ./models \
  --data_dir ./documents \
  --output embeddings.safetensors \
  --file_patterns "*.txt" "*.md" "*.py" \
  --max_files 1000 \
  --compute_similarity \
  --batch_size 16

Output:

• embeddings.safetensors - Compressed embeddings, file info, and similarity matrix

Inference: inference_blt.py

Perform various inference tasks with a trained model:

Next-Byte Prediction

python inference_blt.py \
  --model_dir ./models \
  --mode predict \
  --query "Hello wor" \
  --num_predictions 10

Patch Analysis

python inference_blt.py \
  --model_dir ./models \
  --mode patches \
  --input_file document.txt \
  --output patch_visualization.png

Text Similarity

python inference_blt.py \
  --model_dir ./models \
  --mode similarity \
  --query "Machine learning algorithms" \
  --text2 "Deep learning neural networks"

Similarity Search

python inference_blt.py \
  --model_dir ./models \
  --mode search \
  --query "Python programming tutorial" \
  --embeddings_file embeddings.safetensors \
  --top_k 5