DOCSEARCHbyAlgolia
Submit Docs

retro-pytorch

repository·main·Indexed 21 days ago

https://github.com/lucidrains/retro-pytorch

A PyTorch implementation of DeepMind's Retrieval-based Attention Net (RETRO). It enables high performance with fewer parameters by retrieving information from an external database during inference. The library includes the RETRO class for forward passes, a TrainingWrapper for automating document processing and training, and utilities for FAISS indexing, BERT embedding generation, and precalculating nearest neighbors.

Tokens
2.8K
Snippets
8
Records
8
Agent score
25%

What's inside retro-pytorch

  1. Use TrainingWrapper to process documents and train RETRO

    main

    The TrainingWrapper automates the pipeline of converting a folder of text documents into memmapped numpy arrays and provides a high-level interface for training and generation.

    Workflow:

    1. Instantiate RETRO.
    2. Initialize TrainingWrapper with your RETRO instance and paths for memmapped files.
    3. Use wrapper.get_dataloader() to get a PyTorch DataLoader.
    4. Use wrapper.get_optimizer() to get an AdamW optimizer with correct settings.
    5. Use wrapper.generate() for top-k sampling or prompt-based generation.

    Environment Variable: To force a reprocessing of the training data, run your script with REPROCESS=1.

    $ REPROCESS=1 python train.py
    import torch
from retro_pytorch import RETRO, TrainingWrapper

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()

wrapper = TrainingWrapper(
    retro = retro,
    knn = 2,
    chunk_size = 64,
    documents_path = './text_folder',
    glob = '**/*.txt',
    chunks_memmap_path = './train.chunks.dat',
    seqs_memmap_path = './train.seq.dat',
    doc_ids_memmap_path = './train.doc_ids.dat',
    max_chunks = 1_000_000,
    max_seqs = 100_000,
    knn_extra_neighbors = 100,
    max_index_memory_usage = '100m',
    current_memory_available = '1G'
)

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

# Training step
seq, retrieved = map(lambda t: t.cuda(), next(train_dl))
loss = retro(seq, retrieved, return_loss = True)
loss.backward()
optim.step()
optim.zero_grad()

# Generation
sampled = wrapper.generate(filter_thres = 0.9, temperature = 1.0)
# Or with a prompt
prompt = torch.randint(0, 1000, (1, 128))
sampled = wrapper.generate(prompt, filter_thres = 0.9, temperature = 1.0)

  3. Create a FAISS index and embeddings from chunks

    main

    Use chunks_to_index_and_embed to convert a memmapped chunks numpy array into a FAISS index and a corresponding embeddings array. This allows you to perform similarity searches to find nearest neighbors for retrieval-augmented training.

    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'
)

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

  4. Use the RETRO class for forward passes

    main

    The RETRO class implements the Retrieval-based Attention Net. To perform a forward pass, you must provide both the input sequence and the retrieved tokens.

    Input Shapes:

    • seq: (batch, max_seq_len + 1). The extra token is used for splitting into input and labels during training.
    • retrieved: (batch, num_chunks, num_retrieved_neighbors, retrieved_chunk_with_continuation). For example, if chunk_size is 64, the last dimension is 128 (64 for the chunk + 64 for the continuation).

    Key Parameters for RETRO initialization:

    • chunk_size: The size of the indexed/retrieved chunk.
    • max_seq_len: Maximum sequence length.
    • use_deepnet: Set to True to enable post-normalization with DeepNet residual scaling, which is recommended for scaling to many layers.
    import torch
from retro_pytorch import RETRO

retro = RETRO(
    chunk_size = 64,
    max_seq_len = 2048,
    enc_dim = 896,
    enc_depth = 2,
    dec_dim = 796,
    dec_depth = 12,
    dec_cross_attn_layers = (3, 6, 9, 12),
    heads = 8,
    dim_head = 64,
    dec_attn_dropout = 0.25,
    dec_ff_dropout = 0.25,
    use_deepnet = True
)

seq = torch.randint(0, 20000, (2, 2048 + 1))
retrieved = torch.randint(0, 20000, (2, 32, 2, 128))

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

  5. Generate BERT embeddings and tokenize text

    main

    The retrieval components use a default sentencepiece tokenizer for the cased version of BERT. You can fetch embeddings using bert_embed, which supports either masked mean pooled representations or the [CLS] token representation.

    from retro_pytorch.retrieval import bert_embed, tokenize

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

# Masked mean pooled representation
embeds = bert_embed(ids) # shape: (batch, 768)

# CLS token representation
embeds_cls = bert_embed(ids, return_cls_repr = True) # shape: (batch, 768)

  6. Use RETRODataset for custom data loading

    main

    If you do not want to use TrainingWrapper, you can use RETRODataset to manually assemble data from existing memmapped numpy arrays. This class requires paths to arrays containing chunks, the index of the first chunk in the sequence, and pre-calculated k-nearest neighbor indices.

    import torch
from torch.utils.data import DataLoader
from retro_pytorch import RETRO, RETRODataset

# Requires existing memmapped files:
# - chunk_memmap_path
# - chunk_nn_memmap_path
# - seq_memmap_path

train_ds = RETRODataset(
    num_sequences = 100,
    num_chunks = 1000,
    num_neighbors = 2,
    chunk_size = 64,
    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))

  7. Create chunks using text_folder_to_chunks_

    main

    Use text_folder_to_chunks_ to process a folder of text into the memmapped files required for training. This function calculates chunk start indices and prepares the data for autoregressive training.

    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',
    chunk_size = 64,
    seq_len = 2048,
    max_chunks = 1_000_000,
    max_seqs = 100_000
)
# Returns: {'chunks': <int>, 'docs': <int>, 'seqs': <int>}

  8. Precalculate nearest neighbors for training

    main

    Use chunks_to_precalculated_knn_ to generate a nearest neighbor file (typically saved as .knn.dat) required for RETRO training. This function can filter out neighbors that belong to the same document to ensure retrieval diversity, provided you supply the document IDs.

    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_
    num_nearest_neighbors = 2,                   # number of nearest neighbors you'd like to use
    num_extra_neighbors = 10                     # fetch 10 extra neighbors to account for document filtering
)

# Result is saved to ./train.chunks.knn.dat