iterative-bert

An iterative refinement BERT encoder based on Tiny Recursive Models.

Installation

pip install iterative-bert

Usage

Load a pre-trained encoder from HuggingFace Hub

from iterative_bert.model import IterativeBert

# Load encoder from HuggingFace Hub
encoder = IterativeBert.from_pretrained("your-username/your-model")

# Run inference
import torch
input_ids = torch.tensor([[101, 2054, 2003, 2023, 102]])  # Example tokens
attention_mask = torch.ones_like(input_ids)
outputs = encoder(input_ids, attention_mask=attention_mask)
hidden_states = outputs.last_hidden_state

Create a new encoder from config

from iterative_bert.model import IterativeBert, IterativeBertConfig

config = IterativeBertConfig(
    vocab_size=30522,
    hidden_size=768,
    num_hidden_layers=1,
    num_attention_heads=12,
    intermediate_size=3072,
    h_cycles=1,
    l_cycles=8,
    use_rope=True,
)

encoder = IterativeBert(config)

Features

• Iterative Refinement: Applies transformer layers multiple times with residual connections
• RoPE Support: Rotary Position Embeddings for better length generalization
• Flash Attention: Optional Flash Attention 2/3 support for efficiency
• HuggingFace Compatible: Works with from_pretrained and save_pretrained

License

Apache 2.0