### Setup CALM with Vision Transformer

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

Example configuration using two specialized augmentation LLMs and a vision transformer.

```python
import torch

# pip install vit-pytorch x-transformers

from vit_pytorch.vit import ViT, Attention
from x_transformers import TransformerWrapper, Encoder, Decoder

anchor_llm = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 16,
        dim_head = 2,
        depth = 12,
        heads = 8
    )
)

augment_llm1 = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Encoder(
        dim = 16,
        dim_head = 2,
        depth = 12,
        heads = 8
    )
)

augment_llm2 = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Encoder(
        dim = 16,
        dim_head = 2,
        depth = 12,
        heads = 8
    )
)

vit = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 256,
    depth = 6,
    heads = 16,
    mlp_dim = 2048
)
```

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

### Initialize and Train CALM

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

Setup the anchor and augmentation LLMs, initialize the CALM wrapper, and perform training or generation.

```python
import torch
from x_transformers import TransformerWrapper, Decoder

augment_llm = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 12,
        heads = 8
    )
)

anchor_llm = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 2,
        heads = 8
    )
)

# import CALM wrapper

from CALM_pytorch import CALM, AugmentParams

calm = CALM(
    anchor_llm,
    augment_llms = AugmentParams(
        model = augment_llm,
        connect_every_num_layers = 4
    )
)

# mock input

seq = torch.randint(0, 20000, (1, 1024))
mask = torch.ones((1, 1024)).bool()
prompt = torch.randint(0, 20000, (1, 256))

# forward for finetuning loss

loss = calm(
    seq,
    mask = mask,
    prompt = prompt
)

loss.backward()

# after much training, prompt the composed model

generated = calm.generate(
    prompt = seq[:, :1],
    seq_len = 1024
)
```

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

### Install CALM Pytorch

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

Install the package from PyPI using pip.

```bash
pip install CALM-pytorch
```

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

### Install CALM-pytorch

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

Install the package via pip.

```bash
$ pip install CALM-pytorch
```

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

### Autoregressive Text Generation with Sampling

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

Generate text autoregressively using the composed CALM model with top-p or top-k sampling strategies. Specify the starting prompt, desired sequence length, and the filtering function with its keyword arguments.

```python
import torch
from CALM_pytorch import CALM, AugmentParams
from CALM_pytorch.sampling_utils import top_p, top_k

# Assume calm is already configured and trained

# Generate with top-p (nucleus) sampling (default)
generated = calm.generate(
    prompt=seq[:, :1],          # Starting token(s)
    seq_len=1024,               # Maximum sequence length
    filter_fn=top_p,            # Sampling strategy
    filter_kwargs=dict(thres=0.9)
)

# Generate with top-k sampling
generated_topk = calm.generate(
    prompt=seq[:, :10],
    seq_len=512,
    filter_fn=top_k,
    filter_kwargs=dict(k=50)
)

# Generate with augmentation prompt
generated_with_context = calm.generate(
    prompt=torch.randint(0, 20000, (1, 256)),  # Prompt for augment LLM
    seq_len=512
)
```

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

### CALM Class: Compose Anchor and Augmentation LLMs

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

Compose an anchor LLM with one or more augmentation LLMs. The anchor and augmentation models remain frozen, with only cross-attention parameters trained. Prepare training data and perform a forward pass to get the loss for fine-tuning.

```python
import torch
from x_transformers import TransformerWrapper, Decoder
from CALM_pytorch import CALM, AugmentParams

# Create anchor LLM (the main model that will generate output)
anchor_llm = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Decoder(
        dim=512,
        depth=12,
        heads=8
    )
)

# Create augmentation LLM (provides additional context/knowledge)
augment_llm = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Decoder(
        dim=512,
        depth=12,
        heads=8
    )
)

# Compose models with CALM wrapper
# connect_every_num_layers=4 means anchor attends to augment every 4th layer
calm = CALM(
    anchor_llm,
    augment_llms=AugmentParams(
        model=augment_llm,
        connect_every_num_layers=4
    )
)

# Prepare training data
seq = torch.randint(0, 20000, (1, 1024))      # Target sequence
mask = torch.ones((1, 1024)).bool()            # Attention mask
prompt = torch.randint(0, 20000, (1, 256))     # Prompt for augment LLM

# Forward pass returns cross-entropy loss for fine-tuning
loss = calm(
    seq,
    mask=mask,
    prompt=prompt
)

loss.backward()
# Only cross-attention parameters receive gradients
```

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

### Initialize and Train CALM

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

Demonstrates the initialization of the CALM model with an anchor LLM and multiple augmentation models, followed by a forward pass and backward propagation.

```python
from CALM_pytorch import CALM, AugmentParams, FineTuner

calm = CALM(
    anchor_llm = anchor_llm,
    augment_llms = (
        AugmentParams(
            model = augment_llm1,
            mask_kwarg = 'mask'
        ),
        AugmentParams(
            model = augment_llm2,
            mask_kwarg = 'mask'
        ),
        AugmentParams(
            model = vit,
            input_shape = (3, 256, 256),
            hidden_position = 'input',
            extract_blocks_fn = lambda vit: [m for m in vit.modules() if isinstance(m, Attention)]
        )
    ),
    attn_kwargs = dict(
        linear_project_context = True,
        pre_rmsnorm = True,
        flash = True
    )
)

seq = torch.randint(0, 20000, (1, 1024))
mask = torch.ones((1, 1024)).bool()

prompt = (
    torch.randint(0, 20000, (1, 256)),
    torch.randint(0, 20000, (1, 256)),
    torch.randn(1, 3, 256, 256)
)

loss = calm(
    seq,
    mask = mask,
    prompt = prompt
)

loss.backward()
```

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

### Fine-tune with Accelerate

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

Use the FineTuner class to train the CALM model with 🤗 Accelerate.

```python
trainer = FineTuner(
    calm = calm,
    dataset = dataset,   # returns a dictionary of input kwargs to calm - dict(seq: Tensor, mask: Tensor, prompt: Tensor). it can also return a Tuple, in which data_kwargs needs to be set to the correct ordered value of kwarg names
    batch_size = 16,
    num_train_steps = 10000,
    learning_rate = 3e-4,
    weight_decay = 1e-2,
    warmup_steps = 1000,
    checkpoint_every = 1000
)

trainer()

# checkpoints of the cross attention parameters will be saved to ./checkpoints every 1000 steps
```

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

### Save Model Checkpoint

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

Manually save a model checkpoint to a file.

```python
trainer.save('my_checkpoint.pt')
```

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

### Distributed Fine-Tuning with FineTuner

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

Utilize the FineTuner class for distributed fine-tuning of the CALM model using Hugging Face Accelerate. It supports automatic checkpointing, learning rate warmup, and gradient accumulation. Ensure your dataset provides data in a dictionary format matching the CALM forward signature.

```python
import torch
from torch.utils.data import Dataset
from CALM_pytorch import CALM, AugmentParams, FineTuner

# Custom dataset that returns training data
class CALMDataset(Dataset):
    def __init__(self, size=10000):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        # Return dict with keys matching CALM forward signature
        return {
            'seq': torch.randint(0, 20000, (1024,)),
            'mask': torch.ones(1024).bool(),
            'prompt': torch.randint(0, 20000, (256,))
        }

# Create dataset
dataset = CALMDataset()

# Initialize trainer
trainer = FineTuner(
    calm=calm,
    dataset=dataset,
    batch_size=16,
    num_train_steps=10000,
    learning_rate=3e-4,
    weight_decay=1e-2,
    warmup_steps=1000,
    checkpoint_every=1000,
    checkpoint_path='./checkpoints',
    grad_accum_steps=4  # Gradient accumulation for larger effective batch
)

# Start training - checkpoints saved to ./checkpoints/checkpoint.{n}.pt
trainer()

# Load from checkpoint
trainer.load('checkpoint.5.pt')
```

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

### Perform Sampling with Filtering Utilities

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

Apply top-p or top-k filtering to logits during autoregressive generation.

```python
import torch
from CALM_pytorch.sampling_utils import sample, top_p, top_k

# Top-p (nucleus) sampling - keeps tokens with cumulative probability <= threshold
logits = torch.randn(1, 20000)
filtered_logits = top_p(logits, thres=0.9)

# Top-k sampling - keeps only top k tokens by probability
filtered_logits_k = top_k(logits, k=50)

# Fractional top-k - keeps top fraction of vocabulary
filtered_logits_frac = top_k(logits, frac_num_tokens=0.1)

# Full sampling function for autoregressive generation
from torch.nn import Module

# Assuming `net` is a language model that returns logits
generated = sample(
    net=model,
    prompts=torch.randint(0, 20000, (1, 10)),  # Starting tokens
    seq_len=256,                                # Max length
    temperature=1.0,                            # Sampling temperature
    filter_fn=top_p,                            # Filtering strategy
    filter_kwargs=dict(thres=0.9),
    pad_id=-1,                                  # Padding token ID
    eos_id=None,                                # Optional EOS token
    output_keep_prompt=True                     # Include prompt in output
)
```

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

### Configure Custom Connectivity

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

Define specific layer connections between the anchor and augmentation models using tuples of layer indices.

```python
calm = CALM(
    anchor_llm = anchor_llm,
    augment_llms = (
        AugmentParams(
            model = augment_llm1,
            connections = (
                (1, 12),  # 1st layer of augment llm1 attended to by 12th layer of anchor llm
                (2, 12),
                (3, 12),
                (4, 12),
            ),
        ),
        AugmentParams(
            model = augment_llm2,
            connections = (
                (6, 1),   # 6th layer of augment llm2 attended to by 1st layer of anchor llm
                (6, 2),
                (12, 12),
            )
        )
    )
)
```

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

### Manage CALM State Dictionary

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

Save and load cross-attention parameters while keeping anchor and augment models frozen.

```python
import torch

# CALM only saves/loads cross-attention parameters
# Frozen anchor and augment LLMs are not included

# Get cross-attention state dict
state = calm.state_dict()

# Save to file
torch.save(state, 'calm_cross_attns.pt')

# Load into new CALM instance (must have same architecture)
calm.load_state_dict(torch.load('calm_cross_attns.pt'))

# Get only trainable parameters (cross-attention layers)
trainable_params = list(calm.parameters())
print(f"Trainable parameters: {sum(p.numel() for p in trainable_params)}")
```

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

### Compose LLMs with Custom Connectivity

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

Instantiate the CALM model with anchor and augment LLMs, defining custom connection layers between them. This is useful for creating complex model architectures where specific layers of augment models should connect to specific layers of the anchor model.

```python
import torch
from CALM_pytorch import CALM, AugmentParams

# Assume anchor_llm, augment_llm1, augment_llm2 are pre-defined models

calm = CALM(
    anchor_llm=anchor_llm,
    augment_llms=(
        AugmentParams(
            model=augment_llm1,
            connections=(
                (1, 12),  # augment1 layer 1 -> anchor layer 12
                (2, 12),
                (3, 12),
                (4, 12),
            ),
        ),
        AugmentParams(
            model=augment_llm2,
            connections=(
                (6, 1),   # augment2 layer 6 -> anchor layer 1
                (6, 2),
                (12, 12),
            )
        )
    )
)

# Forward with separate prompts for each augment LLM
seq = torch.randint(0, 20000, (1, 1024))
prompt = (
    torch.randint(0, 20000, (1, 256)),  # Prompt for augment_llm1
    torch.randint(0, 20000, (1, 256)),  # Prompt for augment_llm2
)

loss = calm(seq, prompt=prompt)
loss.backward()
```

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

### Multiple Augmentation LLMs: Compose with Multiple Models

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

CALM supports composing an anchor LLM with multiple augmentation models, each with independent connectivity configurations. Define anchor and multiple augmentation LLMs.

```python
import torch
from x_transformers import TransformerWrapper, Encoder, Decoder
from CALM_pytorch import CALM, AugmentParams

# Anchor model
anchor_llm = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Decoder(dim=256, depth=12, heads=8)
)

# First augmentation LLM (e.g., code specialist)
augment_llm1 = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Encoder(dim=256, depth=12, heads=8)
)

# Second augmentation LLM (e.g., math specialist)
augment_llm2 = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Encoder(dim=256, depth=12, heads=8)
)

# CALM composition with multiple augmentation models would follow here
```

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

### Integrate Vision Transformer for Multimodal Input

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

Integrate a Vision Transformer (ViT) as an augmentation model for multimodal LLM composition. This allows CALM to process image inputs alongside text. Ensure the ViT model is correctly configured with input shape and extraction functions for its layers.

```python
import torch
from vit_pytorch.vit import ViT, Attention
from x_transformers import TransformerWrapper, Encoder, Decoder
from CALM_pytorch import CALM, AugmentParams

# Anchor LLM
anchor_llm = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Decoder(dim=16, dim_head=2, depth=12, heads=8)
)

# Text augmentation LLMs
augment_llm1 = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Encoder(dim=16, dim_head=2, depth=12, heads=8)
)

augment_llm2 = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Encoder(dim=16, dim_head=2, depth=12, heads=8)
)

# Vision Transformer
vit = ViT(
    image_size=256,
    patch_size=32,
    num_classes=1000,
    dim=256,
    depth=6,
    heads=16,
    mlp_dim=2048
)

# Compose all models
calm = CALM(
    anchor_llm=anchor_llm,
    augment_llms=(
        AugmentParams(model=augment_llm1, mask_kwarg='mask'),
        AugmentParams(model=augment_llm2, mask_kwarg='mask'),
        AugmentParams(
            model=vit,
            input_shape=(3, 256, 256),
            hidden_position='input',
            extract_blocks_fn=lambda vit: [
                m for m in vit.modules() if isinstance(m, Attention)
            ]
        )
    ),
    attn_kwargs=dict(
        linear_project_context=True,
        pre_rmsnorm=True,
        flash=True
    )
)

# Forward with multimodal prompts
seq = torch.randint(0, 20000, (1, 1024))
mask = torch.ones((1, 1024)).bool()

prompt = (
    torch.randint(0, 20000, (1, 256)),  # Text prompt for augment_llm1
    torch.randint(0, 20000, (1, 256)),  # Text prompt for augment_llm2
    torch.randn(1, 3, 256, 256)          # Image for ViT
)

loss = calm(seq, mask=mask, prompt=prompt)
loss.backward()
```

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

### Compose Multiple Augmentation LLMs

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

Pass a list of AugmentParams to explore multiple augmentation models.

```python
calm = CALM(
    anchor_llm = anchor_llm,
    augment_llm = [AugmentParams(augment_llm1), AugmentParams(augment_llm2)] # pass in a list of AugmentParams wrapping model and other hparams specific to that transformer
)
```

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

### Extract Hidden States with ExtractHiddensWrapper

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

Use hooks to extract hidden states from transformer blocks during forward passes. Requires the x-transformers library.

```python
import torch
from x_transformers import TransformerWrapper, Decoder
from CALM_pytorch import ExtractHiddensWrapper

# Create model
model = TransformerWrapper(
    num_tokens=20000,
    max_seq_len=1024,
    attn_layers=Decoder(dim=512, depth=6, heads=8)
)

# Get transformer blocks
blocks = []
for layer in model.attn_layers.layers:
    blocks.append(layer[-1])
blocks = blocks[1::2]  # x-transformers specific extraction

# Wrap model to extract hiddens
wrapped = ExtractHiddensWrapper(
    model=model,
    blocks=blocks,
    hidden_positions='output'  # Extract output hidden states
)

# Forward pass returns list of hidden states from each block
seq = torch.randint(0, 20000, (1, 128))
hiddens = wrapped(seq)  # List of tensors, one per block

print(f"Extracted {len(hiddens)} hidden states")
for i, h in enumerate(hiddens):
    print(f"  Block {i}: shape {h.shape}")
```

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

### CALM Citation

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

BibTeX citation for the LLM Augmented LLMs paper.

```bibtex
@inproceedings{Bansal2024LLMAL,
  title   = {LLM Augmented LLMs: Expanding Capabilities through Composition},
  author  = {Rachit Bansal and Bidisha Samanta and Siddharth Dalmia and Nitish Gupta and Shikhar Vashishth and Sriram Ganapathy and Abhishek Bapna and Prateek Jain and Partha Pratim Talukdar},
  year    = {2024},
  url     = {https://api.semanticscholar.org/CorpusID:266755751}
}
```

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

### AugmentParams Dataclass: Configure Augmentation Models

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

Configure augmentation model parameters, including connectivity patterns, hidden state extraction, and model-specific settings. Use basic usage for automatic layer connectivity or advanced usage for custom connections and non-transformer models like ViT.

```python
from dataclasses import dataclass
from CALM_pytorch import AugmentParams

# Basic usage - automatic layer connectivity
params = AugmentParams(
    model=augment_llm,
    connect_every_num_layers=4  # Attend to every 4th layer
)

# Advanced usage - custom layer connections
# Format: (augment_layer, anchor_layer) - 1-indexed
params_custom = AugmentParams(
    model=augment_llm,
    connections=(
        (1, 12),   # Layer 1 of augment attended by layer 12 of anchor
        (2, 12),
        (3, 12),
        (4, 12),
    ),
    mask_kwarg='mask'  # Kwarg name to pass attention mask to model
)

# For non-transformer models (e.g., ViT)
from vit_pytorch.vit import ViT, Attention

vit_params = AugmentParams(
    model=vit_model,
    input_shape=(3, 256, 256),           # Input tensor shape
    hidden_position='input',              # Extract from input not output
    extract_blocks_fn=lambda vit: [       # Custom block extraction
        m for m in vit.modules() if isinstance(m, Attention)
    ]
)
```

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