### Install RETRO Pytorch

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Install the package via pip.

```bash
pip install retro-pytorch
```

--------------------------------

### Training Loop Example

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Standard training loop for the RETRO model. Ensure data is on CUDA and gradients are zeroed after each step.

```python
for step in range(10000):
    seq, retrieved = map(lambda t: t.cuda(), next(train_dl))
    # seq: (batch, 2049) - includes label token
    # retrieved: (batch, 32, 2, 128) - chunks with continuations

    loss = retro(seq, retrieved, return_loss=True)
    loss.backward()
    optim.step()
    optim.zero_grad()

    if step % 100 == 0:
        print(f'Step {step}, Loss: {loss.item():.4f}')
```

--------------------------------

### Training Loop with RETRO

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Perform a single gradient step for training the RETRO model. This involves getting data from the dataloader, calculating the loss, and updating the optimizer.

```python
# get the dataloader and optimizer (AdamW with all the correct settings)

train_dl = iter(wrapper.get_dataloader(batch_size = 2, shuffle = True))
optim = wrapper.get_optimizer(lr = 3e-4, wd = 0.01)

# now do your training
# ex. one gradient step

seq, retrieved = map(lambda t: t.cuda(), next(train_dl))

# seq       - (2, 2049)         - 1 extra token since split by seq[:, :-1], seq[:, 1:]
# retrieved - (2, 32, 2, 128)   - 128 since chunk + continuation, each 64 tokens

loss = retro(
    seq,
    retrieved,
    return_loss = True
)

# one gradient step

loss.backward()
optim.step()
optim.zero_grad()
```

--------------------------------

### Instantiate RETRO and TrainingWrapper

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Initialize the RETRO model and wrap it with TrainingWrapper for streamlined training. Ensure CUDA is available for model placement.

```python
retro = RETRO(
    max_seq_len = 2048,                      # max sequence length
    enc_dim = 896,                           # encoder model dimension
    enc_depth = 3,                           # encoder depth
    dec_dim = 768,                           # decoder model dimensions
    dec_depth = 12,                          # decoder depth
    dec_cross_attn_layers = (1, 3, 6, 9),    # decoder cross attention layers (with causal chunk cross attention)
    heads = 8,                               # attention heads
    dim_head = 64,                           # dimension per head
    dec_attn_dropout = 0.25,                 # decoder attention dropout
    dec_ff_dropout = 0.25                    # decoder feedforward dropout
).cuda()

wrapper = TrainingWrapper(
    retro = retro,                                 # path to retro instance
    knn = 2,                                       # knn (2 in paper was sufficient)
    chunk_size = 64,                               # chunk size (64 in paper)
    documents_path = './text_folder',              # path to folder of text
    glob = '**/*.txt',                             # text glob
    chunks_memmap_path = './train.chunks.dat',     # path to chunks
    seqs_memmap_path = './train.seq.dat',          # path to sequence data
    doc_ids_memmap_path = './train.doc_ids.dat',   # path to document ids per chunk (used for filtering neighbors belonging to same document)
    max_chunks = 1_000_000,                        # maximum cap to chunks
    max_seqs = 100_000,                            # maximum seqs
    knn_extra_neighbors = 100,                     # num extra neighbors to fetch
    max_index_memory_usage = '100m',
    current_memory_available = '1G'
)
```

--------------------------------

### Initialize and Train RETRO Model

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Instantiate the RETRO model with specific dimensions and layers, then perform a forward and backward pass.

```python
retro = RETRO(
    max_seq_len = 2048,                      # max sequence length
    enc_dim = 896,                           # encoder model dimension
    enc_depth = 3,                           # encoder depth
    dec_dim = 768,                           # decoder model dimensions
    dec_depth = 12,                          # decoder depth
    dec_cross_attn_layers = (1, 3, 6, 9),    # decoder cross attention layers (with causal chunk cross attention)
    heads = 8,                               # attention heads
    dim_head = 64,                           # dimension per head
    dec_attn_dropout = 0.25,                 # decoder attention dropout
    dec_ff_dropout = 0.25                    # decoder feedforward dropout
).cuda()

seq, retrieved = map(lambda t: t.cuda(), next(train_dl))

# seq       - (2, 2049)         - 1 extra token since split by seq[:, :-1], seq[:, 1:]
# retrieved - (2, 32, 2, 128)   - 128 since chunk + continuation, each 64 tokens

loss = retro(
    seq,
    retrieved,
    return_loss = True
)

loss.backward()
```

--------------------------------

### Initialize and Run RETRO Model

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Configure the RETRO architecture and perform a forward pass with retrieved chunks. The model supports both retrieval-augmented training and standard transformer inference.

```python
import torch
from retro_pytorch import RETRO

# Initialize RETRO model with encoder-decoder architecture
retro = RETRO(
    chunk_size = 64,                         # chunk size for indexing and retrieval
    max_seq_len = 2048,                      # maximum sequence length
    enc_dim = 896,                           # encoder model dimension
    enc_depth = 2,                           # encoder depth (layers)
    dec_dim = 768,                           # decoder model dimension
    dec_depth = 12,                          # decoder depth (layers)
    dec_cross_attn_layers = (3, 6, 9, 12),   # decoder layers with cross attention
    heads = 8,                               # number of attention heads
    dim_head = 64,                           # dimension per attention head
    dec_attn_dropout = 0.25,                 # decoder attention dropout
    dec_ff_dropout = 0.25,                   # decoder feedforward dropout
    use_deepnet = True                       # enable DeepNet for scaling to 1000+ layers
)

# Prepare input sequence and retrieved chunks
# seq: (batch, seq_len + 1) - extra token for labels
# retrieved: (batch, num_chunks, num_neighbors, chunk_size * 2)
seq = torch.randint(0, 20000, (2, 2048 + 1))
retrieved = torch.randint(0, 20000, (2, 32, 2, 128))  # 32 chunks, 2 neighbors, 128 tokens (chunk + continuation)

# Forward pass with loss computation
loss = retro(seq, retrieved, return_loss=True)
loss.backward()

# Inference without retrieval (standard transformer mode)
logits = retro.forward_without_retrieval(seq[:, :-1])
# Output: (batch, seq_len, vocab_size)
```

--------------------------------

### Initialize and Use RETRO Model

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Initialize the RETRO model with specified parameters and compute the loss for a given sequence and retrieved tokens. Ensure `use_deepnet` is set to True for scaling to large models.

```python
import torch
from retro_pytorch import RETRO

retro = RETRO(
    chunk_size = 64,                         # the chunk size that is indexed and retrieved (needed for proper relative positions as well as causal chunked cross attention)
    max_seq_len = 2048,                      # max sequence length
    enc_dim = 896,                           # encoder model dim
    enc_depth = 2,                           # encoder depth
    dec_dim = 796,                           # decoder model dim
    dec_depth = 12,                          # decoder depth
    dec_cross_attn_layers = (3, 6, 9, 12),   # decoder cross attention layers (with causal chunk cross attention)
    heads = 8,                               # attention heads
    dim_head = 64,                           # dimension per head
    dec_attn_dropout = 0.25,                 # decoder attention dropout
    dec_ff_dropout = 0.25,                   # decoder feedforward dropout
    use_deepnet = True                       # turn on post-normalization with DeepNet residual scaling and initialization, for scaling to 1000 layers
)

seq = torch.randint(0, 20000, (2, 2048 + 1))      # plus one since it is split into input and labels for training
retrieved = torch.randint(0, 20000, (2, 32, 2, 128)) # retrieved tokens - (batch, num chunks, num retrieved neighbors, retrieved chunk with continuation)

loss = retro(seq, retrieved, return_loss = True)
loss.backward()

# do above for many steps
```

--------------------------------

### Generate Mock Chunk Data

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Create mock chunk data for RETRO training using numpy's random integer generation and save it to a memmapped file.

```python
# generate mock chunk data

save_memmap(
    './train.chunks.dat',
    np.int32(np.random.randint(0, 8192, size = (NUM_CHUNKS, CHUNK_SIZE + 1)))
)
```

--------------------------------

### Import RETRO and TrainingWrapper

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Import the necessary classes, RETRO and TrainingWrapper, from the retro-pytorch library. The TrainingWrapper is used for processing text documents into training data.

```python
import torch
from retro_pytorch import RETRO, TrainingWrapper
```

--------------------------------

### Generate Mock Sequence Data

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Create mock sequence data and save it to a memmapped file for RETRO training.

```python
# generate seq data

save_memmap(
    './train.seq.dat',
    np.int32(np.random.randint(0, 128, size = (NUM_SEQS,)))
)
```

--------------------------------

### Instantiate RETRODataset

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Initialize the RETRODataset using paths to memmapped data files. This dataset class constructs sequences and neighbors from the provided data.

```python
# instantiate dataset class
# which constructs the sequence and neighbors from memmapped chunk and neighbor information

train_ds = RETRODataset(
    num_sequences = NUM_SEQS,
    num_chunks = NUM_CHUNKS,
    num_neighbors = NUM_NEIGHBORS,
    chunk_size = CHUNK_SIZE,
    seq_len = 2048,
    chunk_memmap_path = './train.chunks.dat',
    chunk_nn_memmap_path = './train.chunks.knn.dat',
    seq_memmap_path = './train.seq.dat'
)

train_dl = iter(DataLoader(train_ds, batch_size = 2))
```

--------------------------------

### Configure TrainingWrapper for End-to-End Training

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Use the TrainingWrapper to automate text preprocessing, chunk indexing, and KNN computation. This utility provides a streamlined interface for obtaining dataloaders and optimizers.

```python
import torch
from retro_pytorch import RETRO, TrainingWrapper

# Initialize RETRO model
retro = RETRO(
    max_seq_len = 2048,
    enc_dim = 896,
    enc_depth = 3,
    dec_dim = 768,
    dec_depth = 12,
    dec_cross_attn_layers = (1, 3, 6, 9),
    heads = 8,
    dim_head = 64,
    dec_attn_dropout = 0.25,
    dec_ff_dropout = 0.25
).cuda()

# Create training wrapper - processes text folder automatically
wrapper = TrainingWrapper(
    retro = retro,
    knn = 2,                                       # number of nearest neighbors
    chunk_size = 64,                               # chunk size for retrieval
    documents_path = './text_folder',              # folder containing training documents
    glob = '**/*.txt',                             # glob pattern for text files
    chunks_memmap_path = './train.chunks.dat',     # path for chunk storage
    seqs_memmap_path = './train.seq.dat',          # path for sequence indices
    doc_ids_memmap_path = './train.doc_ids.dat',   # path for document IDs
    max_chunks = 1_000_000,                        # maximum number of chunks
    max_seqs = 100_000,                            # maximum number of sequences
    knn_extra_neighbors = 100,                     # extra neighbors for filtering
    max_index_memory_usage = '100m',               # FAISS index memory limit
    current_memory_available = '1G'                # available system memory
)

# Get dataloader and optimizer
train_dl = iter(wrapper.get_dataloader(batch_size=2, shuffle=True))
optim = wrapper.get_optimizer(lr=3e-4, wd=0.01)
```

--------------------------------

### Custom Data Loading with RETRODataset

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Sets up a PyTorch DataLoader using RETRODataset for custom data loading from memory-mapped files. Ensure mock data files are created before instantiation.

```python
import numpy as np
import torch
from torch.utils.data import DataLoader
from retro_pytorch import RETRO, RETRODataset

# Constants for mock data
NUM_CHUNKS = 1000
CHUNK_SIZE = 64
NUM_SEQS = 100
NUM_NEIGHBORS = 2

def save_memmap(path, tensor):
    f = np.memmap(path, dtype=tensor.dtype, mode='w+', shape=tensor.shape)
    f[:] = tensor
    del f

# Create mock chunk data: (num_chunks, chunk_size + 1)
save_memmap(
    './train.chunks.dat',
    np.int32(np.random.randint(0, 8192, size=(NUM_CHUNKS, CHUNK_SIZE + 1)))
)

# Create KNN indices: (num_chunks, num_neighbors)
save_memmap(
    './train.chunks.knn.dat',
    np.int32(np.random.randint(0, 1000, size=(NUM_CHUNKS, NUM_NEIGHBORS)))
)

# Create sequence start indices: (num_seqs,)
save_memmap(
    './train.seq.dat',
    np.int32(np.random.randint(0, 128, size=(NUM_SEQS,)))
)

# Create dataset from memmapped files
train_ds = RETRODataset(
    num_sequences = NUM_SEQS,
    num_chunks = NUM_CHUNKS,
    num_neighbors = NUM_NEIGHBORS,
    chunk_size = CHUNK_SIZE,
    seq_len = 2048,
    chunk_memmap_path = './train.chunks.dat',
    chunk_nn_memmap_path = './train.chunks.knn.dat',
    seq_memmap_path = './train.seq.dat'
)

# Create dataloader
train_dl = DataLoader(train_ds, batch_size=2, shuffle=True)

# Iterate through batches
for seq, retrieved in train_dl:
    # seq: (batch, 2049) - sequence tokens with label
    # retrieved: (batch, 32, 2, 128) - retrieved chunks with continuations
    print(f'Sequence shape: {seq.shape}')
    print(f'Retrieved shape: {retrieved.shape}')
    break
```

--------------------------------

### Text Generation with RETRO

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Demonstrates autoregressive text generation using the TrainingWrapper. Retrieval is automatically handled at chunk boundaries.

```python
import torch
from retro_pytorch import RETRO, TrainingWrapper

# Assume wrapper is already initialized with training data
# wrapper = TrainingWrapper(...)

# Generate from scratch (starts with SOS token)
sampled = wrapper.generate(
    filter_thres = 0.9,      # top-k filtering threshold
    temperature = 1.0         # sampling temperature
)
# Output: (1, <2049) - terminates early if all EOS tokens

# Generate with a prompt
prompt = torch.randint(0, 1000, (1, 128)).cuda()  # two chunks worth of tokens
sampled = wrapper.generate(
    start = prompt,
    filter_thres = 0.9,
    temperature = 1.0
)
# Retrieval happens automatically at chunk boundaries during generation
# Output: (1, <2049) - generated sequence with retrieval-augmented continuation
```

--------------------------------

### Convert Text Folder to Chunks

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Process a directory of text files into memmapped chunk files for training.

```python
from retro_pytorch.retrieval import text_folder_to_chunks_

stats = text_folder_to_chunks_(
    folder = './text_folder',
    glob = '**/*.txt',
    chunks_memmap_path = './train.chunks.dat',
    seqs_memmap_path = './train.seq.dat',
    doc_ids_memmap_path = './train.doc_ids.dat',  # document ids are needed for filtering out neighbors belonging to same document appropriately during computation of nearest neighbors
    chunk_size = 64,
    seq_len = 2048,
    max_chunks = 1_000_000,
    max_seqs = 100_000
)
```

--------------------------------

### Generate Mock Nearest Neighbors Data

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Create mock nearest neighbor data for each chunk and save it to a memmapped file.

```python
# generate nearest neighbors for each chunk

save_memmap(
    './train.chunks.knn.dat',
    np.int32(np.random.randint(0, 1000, size = (NUM_CHUNKS, NUM_NEIGHBORS)))
)
```

--------------------------------

### Build FAISS Index and Embeddings

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Creates BERT embeddings for chunks and builds a FAISS index for efficient nearest neighbor search. Requires specifying paths for chunk data and the index file, along with memory limits.

```python
from retro_pytorch.retrieval import chunks_to_index_and_embed

# Build FAISS index and embeddings from chunk data
index, embeddings = chunks_to_index_and_embed(
    num_chunks = 1000,                           # total number of chunks
    chunk_size = 64,                             # tokens per chunk
    chunk_memmap_path = './train.chunks.dat',    # path to chunk data
    max_rows_per_file = 500,                     # rows per embedding file
    chunks_to_embeddings_batch_size = 16,        # batch size for BERT
    index_file = 'knn.index',                    # FAISS index filename
    max_index_memory_usage = '100m',             # index memory limit
    current_memory_available = '1G'              # system memory available
)

# Query the index for nearest neighbors
query_vector = embeddings[:1]                    # use first embedding as query
distances, indices = index.search(query_vector, k=2)  # find 2 nearest neighbors

# Get neighbor embeddings
neighbor_embeddings = embeddings[indices]        # (1, 2, 768)
print(f'Nearest neighbor indices: {indices}')
print(f'Distances: {distances}')
```

--------------------------------

### Precalculate KNN for Training

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Generate the nearest neighbor file required for training using precomputed chunks.

```python
from retro_pytorch.retrieval import chunks_to_precalculated_knn_

chunks_to_precalculated_knn_(
    num_chunks = 1000,
    chunk_size = 64,
    chunk_memmap_path = './train.chunks.dat',    # path to main chunks dataset
    doc_ids_memmap_path = './train.doc_ids.dat', # path to document ids created by text_folder_to_chunks_, used for filtering out neighbors that belong to the same document
    num_nearest_neighbors = 2,                   # number of nearest neighbors you'd like to use
    num_extra_neighbors = 10                     # fetch 10 extra neighbors, in the case that fetched neighbors are frequently from same document (filtered out)
)
```

--------------------------------

### Generate Text with RETRO (Conditional)

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Generate text using the RETRO model with a given prompt. The wrapper handles retrieval for the initial chunks.

```python
# or you can generate with a prompt, knn retrieval for initial chunks all taken care of

prompt = torch.randint(0, 1000, (1, 128))  # start with two chunks worth of sequence
sampled = wrapper.generate(prompt, filter_thres = 0.9, temperature = 1.0) # (1, <2049) terminates early if all <eos>
```

--------------------------------

### Save Mock Memmap Data

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Utility function to save numpy arrays to memmapped files, used for creating mock data for RETRO training.

```python
import numpy as np

def save_memmap(path, tensor):
    f = np.memmap(path, dtype = tensor.dtype, mode = 'w+', shape = tensor.shape)
    f[:] = tensor
    del f
```

--------------------------------

### Create FAISS Index for Retrieval

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Convert memmapped chunks into embeddings and build a FAISS index for nearest neighbor search.

```python
from retro_pytorch.retrieval import chunks_to_index_and_embed

index, embeddings = chunks_to_index_and_embed(
    num_chunks = 1000,
    chunk_size = 64,
    chunk_memmap_path = './train.chunks.dat'
)

query_vector = embeddings[:1]                   # use first embedding as query
_, indices = index.search(query_vector, k = 2)  # fetch 2 neighbors, first indices should be self

neighbor_embeddings = embeddings[indices]       # (1, 2, 768)
```

--------------------------------

### Reprocess Training Data

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Force a reprocess of the training data by setting the REPROCESS environment flag to 1 when running the training script.

```bash
$ REPROCESS=1 python train.py
```

--------------------------------

### Configure AdamW Optimizer

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Obtains a pre-configured AdamW optimizer with automatic weight decay separation for transformer training. Parameters with fewer than 2 dimensions (biases, layernorms) are excluded from weight decay. `filter_by_requires_grad` can be set to True to only include parameters that require gradients.

```python
from retro_pytorch.optimizer import get_optimizer

# Get optimizer with automatic weight decay separation
# Parameters with ndim < 2 (biases, layernorms) get no weight decay
optimizer = get_optimizer(
    params = retro.parameters(),
    lr = 3e-4,                          # learning rate
    wd = 0.01,                          # weight decay for eligible params
    filter_by_requires_grad = False     # include all params
)

# Training step
loss = retro(seq, retrieved, return_loss=True)
loss.backward()
optimizer.step()
optimizer.zero_grad()

# Or via TrainingWrapper (recommended)
optimizer = wrapper.get_optimizer(lr=3e-4, wd=0.01)
```

--------------------------------

### Generate Text with RETRO (Unconditional)

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Generate text using the RETRO model without a prompt. The generation process uses filtering and temperature for sampling.

```python
# topk sampling with retrieval at chunk boundaries

sampled = wrapper.generate(filter_thres = 0.9, temperature = 1.0) # (1, <2049) terminates early if all <eos>
```

--------------------------------

### BERT Embedding and Tokenization Utilities

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Provides functions for tokenizing text and generating BERT embeddings for semantic similarity search. Embeddings are computed on GPU if available and are suitable for FAISS indexing.

```python
from retro_pytorch.retrieval import bert_embed, tokenize

# Tokenize text strings using BERT tokenizer
ids = tokenize([
    'hello world',
    'foo bar'
])
# Output: tensor of token IDs with padding

# Get masked mean pooled BERT embeddings (default)
embeds = bert_embed(ids)
# Output: (2, 768) - 768 is BERT hidden dimension

# Get CLS token representation instead
embeds_cls = bert_embed(ids, return_cls_repr=True)
# Output: (2, 768) - CLS token embeddings

# Embeddings are computed on GPU if available
# Returns float32 tensor suitable for FAISS indexing
```

--------------------------------

### Text Folder to Chunks Processing

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Processes a folder of text documents into memory-mapped numpy arrays for efficient training. This function creates chunk, sequence index, and document ID files.

```python
from retro_pytorch.retrieval import text_folder_to_chunks_

# Process folder of text files into chunks
stats = text_folder_to_chunks_(
    folder = './text_folder',                    # path to text documents
    glob = '**/*.txt',                           # glob pattern for files
    chunks_memmap_path = './train.chunks.dat',   # output chunks file
    seqs_memmap_path = './train.seq.dat',        # output sequence indices
    doc_ids_memmap_path = './train.doc_ids.dat', # output document IDs per chunk
    chunk_size = 64,                             # tokens per chunk
    seq_len = 2048,                              # sequence length for training
    max_chunks = 1_000_000,                      # maximum chunks to process
    max_seqs = 100_000                           # maximum sequences
)

# Returns statistics dictionary
print(stats)
# {'chunks': 50000, 'docs': 1000, 'seqs': 25000}
# chunks: total number of chunks created
# docs: number of documents processed
# seqs: number of sequences created
```

--------------------------------

### Pre-calculate KNN for Training

Source: https://context7.com/lucidrains/retro-pytorch/llms.txt

Computes and stores k-nearest neighbors for each chunk, filtering out neighbors from the same document. Specify chunk data, document IDs, and desired number of neighbors. Set `force_reprocess` to True to recompute if data exists.

```python
from retro_pytorch.retrieval import chunks_to_precalculated_knn_

# Pre-calculate KNN indices for all chunks
knn_path, faiss_index = chunks_to_precalculated_knn_(
    num_chunks = 1000,                           # total chunks
    chunk_size = 64,                             # tokens per chunk
    chunk_memmap_path = './train.chunks.dat',    # chunk data path
    doc_ids_memmap_path = './train.doc_ids.dat', # document IDs for filtering
    num_nearest_neighbors = 2,                   # KNN to keep per chunk
    num_extra_neighbors = 10,                    # extra neighbors for filtering
    index_file = 'knn.index',                    # FAISS index file
    force_reprocess = False,                     # skip if already processed
    max_index_memory_usage = '100m',
    current_memory_available = '1G'
)

# KNN saved to ./train.chunks.knn.dat
# Each chunk has indices of its 2 nearest neighbors (excluding same document)

print(f'KNN path: {knn_path}')
# Use returned faiss_index for inference-time retrieval
```

--------------------------------

### Generate BERT Embeddings

Source: https://github.com/lucidrains/retro-pytorch/blob/main/README.md

Generate embeddings using BERT, supporting either masked mean pooling or CLS token representation.

```python
from retro_pytorch.retrieval import bert_embed, tokenize

ids = tokenize([
    'hello world',
    'foo bar'
])

embeds = bert_embed(ids) # (2, 768) - 768 is hidden dimension of BERT
```

```python
from retro_pytorch.retrieval import bert_embed, tokenize

ids = tokenize([
    'hello world',
    'foo bar'
])

embeds = bert_embed(ids, return_cls_repr = True) # (2, 768)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.