### Install Enformer PyTorch Locally

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

After cloning, navigate into the directory and install the library using pip. This command installs the package and its dependencies.

```bash
!cd enformer-pytorch && pip install .
```

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

### Setup and Download FASTA Files

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Creates a directory for FASTA files and downloads human and mouse reference genomes from a Google Cloud Storage bucket if they don't already exist. It also decompresses the downloaded .gz files.

```python
import os
fasta_dir = "/root/data/"
!mkdir -p {fasta_dir}
human_fasta_f = 'hg38.ml.fa.gz'
mouse_fasta_f = 'mm10.ml.fa.gz'
human_fasta_gz_path = f"{fasta_dir}/{human_fasta_f}"
mouse_fasta_gz_path = f"{fasta_dir}/{mouse_fasta_f}"
human_fasta_path = human_fasta_gz_path.rstrip(".gz")
mouse_fasta_path = mouse_fasta_gz_path.rstrip(".gz")

if not os.path.isfile(human_fasta_path):
  !gsutil -m cp -n gs://basenji_barnyard/{human_fasta_f} {human_fasta_gz_path}
  !gunzip {human_fasta_gz_path}
if not os.path.isfile(mouse_fasta_path):
  !gsutil -m cp -n gs://basenji_barnyard/{mouse_fasta_f} {mouse_fasta_gz_path}
  !gunzip {mouse_fasta_gz_path}
```

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

### Install enformer-pytorch

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

Install the library using pip.

```bash
pip install enformer-pytorch
```

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

### Install enformer-pytorch

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

Install the library using pip. Ensure you have version 0.5 or higher for full pretrained model compatibility.

```bash
$
pip install enformer-pytorch

```

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

### Install Dependencies with Pip

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Installs necessary Python packages including torchmetrics, kipoiseq, and BioPython. The output is suppressed using '> /dev/null'.

```python
!pip install torchmetrics kipoiseq==0.5.2 BioPython --quiet > /dev/null
```

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

### Download Progress Indicators

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Console output showing download progress for model weights.

```text
Downloading:   0%|          | 0.00/464 [00:00<?, ?B/s]
```

```text
Downloading:   0%|          | 0.00/959M [00:00<?, ?B/s]
```

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

### Initialize Enformer Model

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Load the pre-trained Enformer model and move it to the GPU.

```python
import torch
from enformer_pytorch import Enformer

model = Enformer.from_pretrained("EleutherAI/enformer-official-rough")
```

```python
model = model.eval().cuda()
```

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

### Create and Train Enformer Model

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

Initializes a dataset, dataloader, and Enformer model, followed by a basic training loop.

```python
# Create dataset
ds = GenomeIntervalDataset(
    bed_file='./sequences.bed',           # BED file with genomic intervals
    fasta_file='./hg38.ml.fa',            # FASTA file with genome sequence
    filter_df_fn=filter_train,            # Filter function for dataframe
    return_seq_indices=True,              # Return indices instead of one-hot
    shift_augs=(-2, 2),                   # Random shift augmentation range
    context_length=196_608,               # Context length (padded if needed)
    chr_bed_to_fasta_map={                # Optional chromosome name mapping
        'chr1': 'chromosome1',
        'chr2': 'chromosome2',
    }
)

# Create DataLoader
dataloader = DataLoader(ds, batch_size=4, shuffle=True)

# Create model
model = Enformer.from_hparams(
    dim=1536, depth=11, heads=8,
    output_heads=dict(human=5313, mouse=1643),
    target_length=896
)

# Training loop
for seq in dataloader:
    pred = model(seq, head='human')
    print(pred.shape)  # torch.Size([4, 896, 5313])
    break
```

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

### Instantiate Enformer Model

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

Initialize the Enformer model with specified hyperparameters. The model can predict outputs for 'human' and 'mouse' based on the provided output heads.

```python
import torch
from enformer_pytorch import Enformer

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = dict(human = 5313, mouse = 1643),
    target_length = 896,
)
    
seq = torch.randint(0, 5, (1, 196_608)) # for ACGTN, in that order (-1 for padding)
output = model(seq)

output['human'] # (1, 896, 5313)
output['mouse'] # (1, 896, 1643)

```

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

### Import Libraries and Define Constants

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Imports core libraries like torch, numpy, tensorflow, and pandas. Defines constants for sequence and target lengths, and sets up paths for human and mouse FASTA files.

```python
import torch
import numpy as np
import tensorflow as tf
import os 
import json
import pandas as pd
import pyfaidx
import kipoiseq
import functools
from kipoiseq import Interval

SEQUENCE_LENGTH = 196_608
BIN_SIZE = 128
TARGET_LENGTH = 896
```

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

### Initialize Enformer with from_hparams

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

Configures the Enformer model architecture using specific hyperparameters. Enable use_checkpointing for memory efficiency and use_tf_gamma for compatibility with pretrained weights.

```python
model = Enformer.from_hparams(
    dim=1536,
    depth=11,
    heads=8,
    output_heads=dict(human=5313, mouse=1643),
    target_length=896,
    use_checkpointing=True,  # Enable for memory efficiency
    use_tf_gamma=True,       # For compatibility with pretrained weights
)
```

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

### Fine-tuning with HeadAdapterWrapper

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

Add a custom head adapter to the Enformer model for new genomic tracks, with options to freeze the base model or specific layers.

```python
import torch
from enformer_pytorch import from_pretrained
from enformer_pytorch.finetune import HeadAdapterWrapper

# Load pretrained Enformer
enformer = from_pretrained('EleutherAI/enformer-official-rough')

# Wrap with head adapter for new tracks
model = HeadAdapterWrapper(
    enformer=enformer,
    num_tracks=128,                   # Number of new output tracks
    post_transformer_embed=False,     # Use final embeddings (after pointwise conv)
).cuda()

# Prepare data
seq = torch.randint(0, 5, (1, 196_608 // 2)).cuda()
target = torch.randn(1, 200, 128).cuda()  # 128 new tracks

# Training with frozen Enformer
loss = model(seq, target=target, freeze_enformer=True)
loss.backward()
print(f"Loss (frozen): {loss.item():.4f}")

# Training with only LayerNorms unfrozen
loss = model(seq, target=target, finetune_enformer_ln_only=True)
loss.backward()
print(f"Loss (LN only): {loss.item():.4f}")

# Training with last N transformer layers unfrozen
loss = model(seq, target=target, finetune_last_n_layers_only=3)
loss.backward()
print(f"Loss (last 3 layers): {loss.item():.4f}")
```

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

### Clone Enformer PyTorch Repository

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Use this command to download the Enformer PyTorch source code from GitHub.

```bash
!git clone https://github.com/lucidrains/enformer-pytorch.git
```

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

### Initialize Enformer Model with Custom Hyperparameters

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

Create an Enformer model instance with specified dimensions, depth, attention heads, output tracks, and target length. Input DNA sequences should be integer indices.

```python
import torch
from enformer_pytorch import Enformer

# Create Enformer model with custom hyperparameters
model = Enformer.from_hparams(
    dim=1536,                                    # Model dimension
    depth=11,                                    # Number of transformer layers
    heads=8,                                     # Number of attention heads
    output_heads=dict(human=5313, mouse=1643),  # Output tracks per species
    target_length=896,                           # Target output length
)

# Input: DNA sequence as integer indices (0=A, 1=C, 2=G, 3=T, 4=N, -1=padding)
seq = torch.randint(0, 5, (1, 196_608))

# Forward pass returns predictions for all output heads
output = model(seq)

print(output['human'].shape)  # torch.Size([1, 896, 5313])
print(output['mouse'].shape)  # torch.Size([1, 896, 1643])
```

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

### Load Pretrained Enformer Model

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

Load the official Enformer model from Hugging Face. This version may have minor numerical discrepancies but generally performs well.

```python
from enformer_pytorch import from_pretrained

enformer = from_pretrained('EleutherAI/enformer-official-rough')

```

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

### Fine-tune Enformer with HeadAdapterWrapper

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

Adapt the Enformer model for fine-tuning on new tracks by wrapping it with `HeadAdapterWrapper`. This wrapper allows for adjustments to embeddings taken after the final pointwise block.

```python
import torch
from enformer_pytorch import from_pretrained
from enformer_pytorch.finetune import HeadAdapterWrapper

enformer = from_pretrained('EleutherAI/enformer-official-rough')

model = HeadAdapterWrapper(
    enformer = enformer,
    num_tracks = 128,
    post_transformer_embed = False   # by default, embeddings are taken from after the final pointwise block w/ conv -> gelu - but if you'd like the embeddings right after the transformer block with a learned layernorm, set this to True
).cuda()

seq = torch.randint(0, 5, (1, 196_608 // 2,)).cuda()
target = torch.randn(1, 200, 128).cuda()  # 128 tracks

loss = model(seq, target = target)
loss.backward()
```

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

### Fine-tune Enformer with ContextAttentionAdapterWrapper

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

Utilize `ContextAttentionAdapterWrapper` for fine-tuning by employing attention aggregation from a set of context embeddings. Configure the number of heads and dimension per head for the cross-attention mechanism.

```python
import torch
from enformer_pytorch import from_pretrained
from enformer_pytorch.finetune import ContextAttentionAdapterWrapper

enformer = from_pretrained('EleutherAI/enformer-official-rough')
    
model = ContextAttentionAdapterWrapper(
    enformer = enformer,
    context_dim = 1024,
    heads = 8,              # number of heads in the cross attention
    dim_head = 64           # dimension per head
).cuda()

seq = torch.randint(0, 5, (1, 196_608 // 2,)).cuda()

target = torch.randn(1, 200, 4).cuda()      # 4 tracks
context = torch.randn(4, 16, 1024).cuda()   # 4 contexts for the different 'tracks', each with 16 tokens

context_mask = torch.ones(4, 16).bool().cuda() # optional context mask, in example, include all context tokens

loss = model(
    seq,
    context = context,
    context_mask = context_mask,
    target = target
)

loss.backward()
```

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

### Enable Checkpointing for Memory Saving

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

Use `use_checkpointing = True` when loading a pre-trained Enformer model to reduce memory consumption during fine-tuning.

```python
from enformer_pytorch import from_pretrained

enformer = from_pretrained('EleutherAI/enformer-official-rough', use_checkpointing = True)

# finetune enformer on a limited budget
```

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

### Fine-tune Enformer with ContextAdapterWrapper

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

Fine-tune Enformer on contextual data (e.g., cell type, transcription factor) using the `ContextAdapterWrapper`. This wrapper integrates context dimensions into the model.

```python
import torch
from enformer_pytorch import from_pretrained
from enformer_pytorch.finetune import ContextAdapterWrapper

enformer = from_pretrained('EleutherAI/enformer-official-rough')
    
model = ContextAdapterWrapper(
    enformer = enformer,
    context_dim = 1024
).cuda()

seq = torch.randint(0, 5, (1, 196_608 // 2,)).cuda()

target = torch.randn(1, 200, 4).cuda()  # 4 tracks
context = torch.randn(4, 1024).cuda()   # 4 contexts for the different 'tracks'

loss = model(
    seq,
    context = context,
    target = target
)

loss.backward()
```

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

### Load Pretrained Enformer Weights

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

Load official DeepMind Enformer weights from HuggingFace. Options include loading with custom target length, dropout rate, or enabling gradient checkpointing for memory efficiency.

```python
from enformer_pytorch import from_pretrained

# Load pretrained model from HuggingFace
enformer = from_pretrained('EleutherAI/enformer-official-rough')

# Load with custom target length for fine-tuning on shorter sequences
model = from_pretrained(
    'EleutherAI/enformer-official-rough',
    target_length=128,
    dropout_rate=0.1
)

# Enable gradient checkpointing for memory-efficient fine-tuning
enformer_memory_efficient = from_pretrained(
    'EleutherAI/enformer-official-rough',
    use_checkpointing=True
)

# Run inference
seq = torch.randint(0, 5, (1, 196_608))
output = enformer(seq)
print(output['human'].shape)  # torch.Size([1, 896, 5313])
```

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

### Download Basenji Data

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Shell commands to download the required FASTA files for human and mouse genomes.

```bash
Copying gs://basenji_barnyard/hg38.ml.fa.gz...
/ [1/1 files][839.8 MiB/839.8 MiB] 100% Done  58.5 MiB/s ETA 00:00:00           
Operation completed over 1 objects/839.8 MiB.                                    
Copying gs://basenji_barnyard/mm10.ml.fa.gz...
/ [1/1 files][800.8 MiB/800.8 MiB] 100% Done  65.9 MiB/s ETA 00:00:00           
Operation completed over 1 objects/800.8 MiB.                                    
```

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

### Load Pretrained Enformer with Custom Target Length

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

Load the pretrained Enformer model, overriding the default target length and dropout rate for specific use cases with shorter sequences.

```python
from enformer_pytorch import from_pretrained

model = from_pretrained('EleutherAI/enformer-official-rough', target_length = 128, dropout_rate = 0.1)

```

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

### Enformer for Training with Poisson Loss

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

Calculate the Poisson loss for training by providing the sequence, target head, and target data. The loss can then be backpropagated.

```python
import torch
from enformer_pytorch import Enformer, seq_indices_to_one_hot

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = dict(human = 5313, mouse = 1643),
    target_length = 200,
).cuda()

seq = torch.randint(0, 5, (196_608 // 2,)).cuda()
target = torch.randn(200, 5313).cuda()

loss = model(
    seq,
    head = 'human',
    target = target
)

loss.backward()

# after much training

corr_coef = model(
    seq,
    head = 'human',
    target = target,
    return_corr_coef = True
)

corr_coef # pearson R, used as a metric in the paper

```

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

### Attention-Based Context Aggregation

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

Use cross-attention to aggregate multiple context tokens per track for more expressive context conditioning.

```python
import torch
from enformer_pytorch import from_pretrained
from enformer_pytorch.finetune import ContextAttentionAdapterWrapper

# Load pretrained model
enformer = from_pretrained('EleutherAI/enformer-official-rough')

# Create attention-based context adapter
model = ContextAttentionAdapterWrapper(
    enformer=enformer,
    context_dim=1024,
    heads=8,           # Cross-attention heads
    dim_head=64,       # Dimension per head
).cuda()

# Prepare data with multiple context tokens per track
seq = torch.randint(0, 5, (1, 196_608 // 2)).cuda()
target = torch.randn(1, 200, 4).cuda()          # 4 tracks
context = torch.randn(4, 16, 1024).cuda()       # 4 contexts, each with 16 tokens
context_mask = torch.ones(4, 16).bool().cuda()  # Optional mask for context tokens

# Forward pass
loss = model(
    seq,
    context=context,
    context_mask=context_mask,
    target=target,
    freeze_enformer=True
)

loss.backward()
print(f"Attention context loss: {loss.item():.4f}")

# Inference
predictions = model(seq, context=context, context_mask=context_mask)
print(predictions.shape)  # torch.Size([1, 200, 4])
```

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

### Context-Conditioned Predictions with ContextAdapterWrapper

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

Fine-tune the model using contextual information like cell type or transcription factor embeddings via a hypernetwork approach.

```python
import torch
from enformer_pytorch import from_pretrained
from enformer_pytorch.finetune import ContextAdapterWrapper

# Load pretrained model
enformer = from_pretrained('EleutherAI/enformer-official-rough')

# Create context-conditioned model
model = ContextAdapterWrapper(
    enformer=enformer,
    context_dim=1024,  # Dimension of context embeddings
).cuda()

# Prepare data
seq = torch.randint(0, 5, (1, 196_608 // 2)).cuda()
target = torch.randn(1, 200, 4).cuda()   # 4 tracks
context = torch.randn(4, 1024).cuda()    # 4 context embeddings (one per track)

# Forward pass with context
loss = model(
    seq,
    context=context,
    target=target,
    freeze_enformer=True
)

loss.backward()
print(f"Context-conditioned loss: {loss.item():.4f}")

# Inference without target
predictions = model(seq, context=context)
print(predictions.shape)  # torch.Size([1, 200, 4])
```

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

### Evaluation Progress Output

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Console output showing the progress of the correlation computation.

```text
100%|██████████| 100/100 [01:22<00:00,  1.21it/s]
```

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

### EnformerConfig for Custom Configurations

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

Defines custom hyperparameters for the Enformer model using the EnformerConfig class.

```python
from enformer_pytorch import Enformer
from enformer_pytorch.config_enformer import EnformerConfig

# Create custom configuration
config = EnformerConfig(
    dim=1536,                    # Model dimension
    depth=11,                    # Transformer depth
    heads=8,                     # Attention heads
    output_heads=dict(           # Custom output heads
        human=5313,
        mouse=1643,
        custom_track=100
    ),
    target_length=896,           # Output sequence length
    attn_dim_key=64,             # Attention key dimension
    dropout_rate=0.4,            # Dropout rate
    attn_dropout=0.05,           # Attention dropout
    pos_dropout=0.01,            # Positional encoding dropout
    use_checkpointing=False,     # Gradient checkpointing
    num_downsamples=7,           # Number of downsampling layers (128x total)
    dim_divisible_by=128,        # Dimension divisibility constraint
    use_tf_gamma=False,          # Use TensorFlow gamma positions
)

# Create model from config
model = Enformer(config)
```

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

### MeanPearsonCorrCoefPerChannel Metric

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

Computes the mean Pearson correlation coefficient across channels, compatible with distributed training.

```python
import torch
from enformer_pytorch.metrics import MeanPearsonCorrCoefPerChannel

# Initialize metric for 5313 channels (human tracks)
metric = MeanPearsonCorrCoefPerChannel(n_channels=5313)

# Simulate predictions and targets over multiple batches
for batch in range(10):
    preds = torch.randn(4, 896, 5313)   # Batch of predictions
    target = torch.randn(4, 896, 5313)  # Batch of targets
    metric.update(preds, target)

# Compute final correlation across all batches
correlation = metric.compute()
print(correlation.shape)              # torch.Size([5313])
print(f"Mean correlation: {correlation.mean().item():.4f}")

# Reset for new epoch
metric.reset()
```

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

### Load Genomic Sequences with Enformer

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

Use GenomicIntervalDataset to load sequences from a BED file. Configure FASTA file path, filtering functions, sequence representation (indices or one-hot), and context length. Optional chromosome name mapping is supported.

```python
import torch
import polars as pl
from enformer_pytorch import Enformer, GenomeIntervalDataset

filter_train = lambda df: df.filter(pl.col('column_4') == 'train')

ds = GenomeIntervalDataset(
    bed_file = './sequences.bed',                       # bed file - columns 0, 1, 2 must be <chromosome>, <start position>, <end position>
    fasta_file = './hg38.ml.fa',                        # path to fasta file
    filter_df_fn = filter_train,                        # filter dataframe function
    return_seq_indices = True,                          # return nucleotide indices (ACGTN) or one hot encodings
    shift_augs = (-2, 2),                               # random shift augmentations from -2 to +2 basepairs
    context_length = 196_608,
    # this can be longer than the interval designated in the .bed file,
    # in which case it will take care of lengthening the interval on either sides
    # as well as proper padding if at the end of the chromosomes
    chr_bed_to_fasta_map = {
        'chr1': 'chromosome1',  # if the chromosome name in the .bed file is different than the key name in the fasta file, you can rename them on the fly
        'chr2': 'chromosome2',
        'chr3': 'chromosome3',
        # etc etc
    }
)

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = dict(human = 5313, mouse = 1643),
    target_length = 896,
)

seq = ds[0] # (196608,)
pred = model(seq, head = 'human') # (896, 5313)
```

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

### BasenjiDataSet Class - Static Methods

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Provides static methods for the Basenji dataset, including retrieving organism paths, loading metadata from JSON, and performing one-hot encoding of DNA sequences.

```python
class BasenjiDataSet(torch.utils.data.IterableDataset):
  @staticmethod
  def get_organism_path(organism):
    return os.path.join('gs://basenji_barnyard/data', organism)
  @classmethod
  def get_metadata(cls, organism):
    # Keys:
    # num_targets, train_seqs, valid_seqs, test_seqs, seq_length,
    # pool_width, crop_bp, target_length
    path = os.path.join(cls.get_organism_path(organism), 'statistics.json')
    with tf.io.gfile.GFile(path, 'r') as f:
      return json.load(f)
  @staticmethod
  def one_hot_encode(sequence):
    return kipoiseq.transforms.functional.one_hot_dna(sequence).astype(np.float32)

  @classmethod
  def get_tfrecord_files(cls, organism, subset):
    # Sort the values by int(*).
    return sorted(tf.io.gfile.glob(os.path.join(
        cls.get_organism_path(organism), 'tfrecords', f'{subset}-*.tfr'
      )), key=lambda x: int(x.split('-')[-1].split('.')[0]))
  
  @property
  def num_channels(self):
    metadata = self.get_metadata(self.organism)
    return metadata['num_targets']

  @staticmethod
  def deserialize(serialized_example, metadata):
    """Deserialize bytes stored in TFRecordFile."""
    # Deserialization
    feature_map = {
        'sequence': tf.io.FixedLenFeature([], tf.string),  # Ignore this, resize our own bigger one
        'target': tf.io.FixedLenFeature([], tf.string),
    }
    example = tf.io.parse_example(serialized_example, feature_map)
    sequence = tf.io.decode_raw(example['sequence'], tf.bool)

```

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

### Deserialize and Load Basenji Data

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Methods for deserializing TFRecord data and initializing the dataset with sequence and target processing.

```python
    sequence = tf.reshape(sequence, (metadata['seq_length'], 4))
    sequence = tf.cast(sequence, tf.float32)

    target = tf.io.decode_raw(example['target'], tf.float16)
    target = tf.reshape(target,
                        (metadata['target_length'], metadata['num_targets']))
    target = tf.cast(target, tf.float32)

    return {'sequence_old': sequence,
            'target': target}
```

```python
  @classmethod
  def get_dataset(cls, organism, subset, num_threads=8):
    metadata = cls.get_metadata(organism)
    dataset = tf.data.TFRecordDataset(cls.get_tfrecord_files(organism, subset),
                                      compression_type='ZLIB',
                                      num_parallel_reads=num_threads).map(
                                          functools.partial(cls.deserialize, metadata=metadata)
                                      )
    return dataset
```

```python
  def __init__(self, organism:str, subset:str, seq_len:int, fasta_path:str, n_to_test:int = -1):
    assert subset in {"train", "valid", "test"}
    assert organism in {"human", "mouse"}
    self.organism = organism
    self.subset = subset
    self.base_dir = self.get_organism_path(organism)
    self.seq_len = seq_len
    self.fasta_reader = FastaStringExtractor(fasta_path)
    self.n_to_test = n_to_test
    with tf.io.gfile.GFile(f"{self.base_dir}/sequences.bed", 'r') as f:
      region_df = pd.read_csv(f, sep="\t", header=None)
      region_df.columns = ['chrom', 'start', 'end', 'subset']
      self.region_df = region_df.query('subset==@subset').reset_index(drop=True)
```

```python
  def __iter__(self):
    worker_info = torch.utils.data.get_worker_info()
    assert worker_info is None, "Only support single process loading"
    # If num_threads > 1, the following will actually shuffle the inputs! luckily we catch this with the sequence comparison
    basenji_iterator = self.get_dataset(self.organism, self.subset, num_threads=1).as_numpy_iterator()
    for i, records in enumerate(basenji_iterator):
      loc_row = self.region_df.iloc[i]
      target_interval = Interval(loc_row['chrom'], loc_row['start'], loc_row['end'])
      sequence_one_hot = self.one_hot_encode(self.fasta_reader.extract(target_interval.resize(self.seq_len)))
      if self.n_to_test >= 0 and i < self.n_to_test:
        old_sequence_onehot = records["sequence_old"]
        if old_sequence_onehot.shape[0] > sequence_one_hot.shape[0]:
          diff = old_sequence_onehot.shape[0] - sequence_one_hot.shape[0]
          trim = diff//2
          np.testing.assert_equal(old_sequence_onehot[trim:(-trim)], sequence_one_hot)
        elif sequence_one_hot.shape[0] > old_sequence_onehot.shape[0]:
          diff = sequence_one_hot.shape[0] - old_sequence_onehot.shape[0]
          trim = diff//2
          np.testing.assert_equal(old_sequence_onehot, sequence_one_hot[trim:(-trim)])
        else:
          np.testing.assert_equal(old_sequence_onehot, sequence_one_hot)
      yield {
          "sequence": sequence_one_hot,
          "target": records["target"],
      }
```

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

### Evaluate Model Correlation

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Compute Pearson correlation coefficients per channel for the model predictions.

```python
from enformer_pytorch.metrics import MeanPearsonCorrCoefPerChannel
```

```python
from tqdm import tqdm
from torchmetrics.regression.pearson import PearsonCorrCoef
def compute_correlation(model, organism:str="human", subset:str="valid", max_steps=-1):
  fasta_path = human_fasta_path if organism == "human" else mouse_fasta_path
  ds = BasenjiDataSet(organism, subset, SEQUENCE_LENGTH, fasta_path)
  total = len(ds.region_df) # number of records
  dl = torch.utils.data.DataLoader(ds, num_workers=0, batch_size=1)
  corr_coef = MeanPearsonCorrCoefPerChannel(n_channels=ds.num_channels)
  n_steps = total if max_steps <= 0 else max_steps
  for i,batch in enumerate(tqdm(dl, total=n_steps)):
    if max_steps > 0 and i >= max_steps:
      break
    batch_gpu = {k:v.to(model.device) for k,v in batch.items()}
    sequence = batch_gpu['sequence']
    target = batch_gpu['target']
    with torch.no_grad():
      pred = model(sequence)[organism]
      corr_coef(preds=pred.cpu(), target=target.cpu())
  return corr_coef.compute().mean()
compute_correlation(model, organism="human", subset="valid", max_steps=100)
```

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

### Compute Correlation for Data Subsets

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

Calculates the correlation for specific organism data subsets. The max_steps parameter controls the number of steps for the computation.

```python
compute_correlation(model, organism="human", subset="valid", max_steps=-1)
```

```python
compute_correlation(model, organism="human", subset="test", max_steps=-1)
```

```python
compute_correlation(model, organism="human", subset="train", max_steps=-1)
```

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

### FastaInterval for Direct Sequence Access

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

Accesses genomic sequences directly from FASTA files with support for padding and augmentation.

```python
import torch
from enformer_pytorch import FastaInterval

# Create FASTA interval accessor
fasta = FastaInterval(
    fasta_file='./hg38.ml.fa',
    context_length=196_608,
    return_seq_indices=True,
    shift_augs=(-2, 2),
    rc_aug=True
)

# Get sequence for a specific interval
# Automatically pads if interval extends beyond chromosome boundaries
seq = fasta(
    chr_name='chr1',
    start=100000,
    end=100500,
    return_augs=False
)
print(seq.shape)  # torch.Size([196608])

# Get sequence with augmentation metadata
seq, shift, rc = fasta('chr1', 100000, 100500, return_augs=True)
print(f"Shift: {shift.item()}, Reverse complement: {rc.item()}")
```

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

### Convert DNA Sequences to One-Hot Encoding

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

Utilities to convert DNA sequences from integer indices or string format to one-hot encodings suitable for model input. Supports batch conversion.

```python
import torch
from enformer_pytorch import seq_indices_to_one_hot, str_to_one_hot, Enformer

# Convert sequence indices to one-hot encoding
seq_indices = torch.randint(0, 5, (1, 196_608))  # 0=A, 1=C, 2=G, 3=T, 4=N
one_hot = seq_indices_to_one_hot(seq_indices)
print(one_hot.shape)  # torch.Size([1, 196608, 4])

# Convert DNA string directly to one-hot
dna_string = "ACGTACGTNNACGT"
one_hot_from_str = str_to_one_hot(dna_string)
print(one_hot_from_str.shape)  # torch.Size([14, 4])

# Batch conversion from strings
batch_strings = ["ACGTACGT", "TGCATGCA"]
one_hot_batch = str_to_one_hot(batch_strings)
print(one_hot_batch.shape)  # torch.Size([2, 8, 4])

# Use one-hot input directly with model
model = Enformer.from_hparams(
    dim=1536, depth=11, heads=8,
    output_heads=dict(human=5313, mouse=1643),
    target_length=896
)
output = model(one_hot)
print(output['human'].shape)  # torch.Size([1, 896, 5313])
```

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

### Loading Genomic Data with GenomeIntervalDataset

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

Load genomic sequences from BED and FASTA files, including support for custom filtering functions.

```python
import torch
import polars as pl
from enformer_pytorch import Enformer, GenomeIntervalDataset
from torch.utils.data import DataLoader

# Define filter function for training data
filter_train = lambda df: df.filter(pl.col('column_4') == 'train')
```

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

### Train Enformer Model with Poisson Loss

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

Train the Enformer model using Poisson loss for gene expression prediction. The model can compute loss directly when target data is provided. Evaluation can be done using Pearson correlation coefficient.

```python
import torch
from enformer_pytorch import Enformer

model = Enformer.from_hparams(
    dim=1536,
    depth=11,
    heads=8,
    output_heads=dict(human=5313, mouse=1643),
    target_length=200,
).cuda()

# Prepare sequence and target data
seq = torch.randint(0, 5, (196_608 // 2,)).cuda()
target = torch.randn(200, 5313).cuda()

# Compute Poisson loss directly
loss = model(
    seq,
    head='human',
    target=target
)

# Backpropagate
loss.backward()
print(f"Training loss: {loss.item():.4f}")

# After training, evaluate with Pearson correlation coefficient
with torch.no_grad():
    corr_coef = model(
        seq,
        head='human',
        target=target,
        return_corr_coef=True
    )
    print(f"Pearson R: {corr_coef.item():.4f}")
```

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

### Enformer with Embeddings Return

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

Retrieve embeddings from the Enformer model by setting `return_embeddings` to `True`. Embeddings can be used for fine-tuning.

```python
import torch
from enformer_pytorch import Enformer, seq_indices_to_one_hot

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = dict(human = 5313, mouse = 1643),
    target_length = 896,
)

seq = torch.randint(0, 5, (1, 196_608))
one_hot = seq_indices_to_one_hot(seq)

output, embeddings = model(one_hot, return_embeddings = True)

embeddings # (1, 896, 3072)

```

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

### GenomeIntervalDataset with Augmentation Metadata

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

Configures a dataset to return augmentation metadata, such as shift values and reverse complement status, for target alignment.

```python
import torch
import polars as pl
from enformer_pytorch import GenomeIntervalDataset

# Create dataset with augmentation metadata
ds = GenomeIntervalDataset(
    bed_file='./sequences.bed',
    fasta_file='./hg38.ml.fa',
    filter_df_fn=lambda df: df.filter(pl.col('column_4') == 'train'),
    return_seq_indices=True,
    shift_augs=(-2, 2),              # Random shift augmentation
    rc_aug=True,                      # Enable reverse complement augmentation
    context_length=196_608,
    return_augs=True                  # Return augmentation metadata
)

# Get sequence with augmentation info
seq, rand_shift_val, rc_bool = ds[0]

print(seq.shape)           # torch.Size([196608])
print(rand_shift_val)      # tensor([1]) - random shift value
print(rc_bool)             # tensor([False]) - whether reverse complement was applied

# Use augmentation info to align target data
# If rc_bool is True, reverse complement the target as well
```

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

### Extracting Model Embeddings

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

Extract intermediate embeddings from the Enformer model for downstream analysis, either alongside predictions or independently.

```python
import torch
from enformer_pytorch import Enformer, seq_indices_to_one_hot

model = Enformer.from_hparams(
    dim=1536,
    depth=11,
    heads=8,
    output_heads=dict(human=5313, mouse=1643),
    target_length=896,
)

seq = torch.randint(0, 5, (1, 196_608))
one_hot = seq_indices_to_one_hot(seq)

# Get predictions along with embeddings
output, embeddings = model(one_hot, return_embeddings=True)
print(embeddings.shape)  # torch.Size([1, 896, 3072])

# Get only embeddings without predictions
embeddings_only = model(one_hot, return_only_embeddings=True)
print(embeddings_only.shape)  # torch.Size([1, 896, 3072])
```

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

### Enformer with One-Hot Encoding

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

Convert sequence indices to one-hot encodings before passing them to the model. This method requires the input sequence to be float values.

```python
import torch
from enformer_pytorch import Enformer, seq_indices_to_one_hot

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = dict(human = 5313, mouse = 1643),
    target_length = 896,
)

seq = torch.randint(0, 5, (1, 196_608))
one_hot = seq_indices_to_one_hot(seq)

output = model(one_hot)

output['human'] # (1, 896, 5313)
output['mouse'] # (1, 896, 1643)

```

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

### Enformer with Reverse Complement Augmentation

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

Configure GenomicIntervalDataset to return augmentation metadata, including random shift values and reverse complement status. This is useful for aligning target data with augmented sequences.

```python
import torch
import polars as pl
from enformer_pytorch import Enformer, GenomeIntervalDataset

filter_train = lambda df: df.filter(pl.col('column_4') == 'train')

ds = GenomeIntervalDataset(
    bed_file = './sequences.bed',                       # bed file - columns 0, 1, 2 must be <chromosome>, <start position>, <end position>
    fasta_file = './hg38.ml.fa',                        # path to fasta file
    filter_df_fn = filter_train,                        # filter dataframe function
    return_seq_indices = True,                          # return nucleotide indices (ACGTN) or one hot encodings
    shift_augs = (-2, 2),                               # random shift augmentations from -2 to +2 basepairs
    rc_aug = True,                                      # use reverse complement augmentation with 50% probability
    context_length = 196_608,
    return_augs = True                                  # return the augmentation meta data
)

seq, rand_shift_val, rc_bool = ds[0] # (196608,), (1,), (1,)
```

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

### FastaStringExtractor Class

Source: https://github.com/lucidrains/enformer-pytorch/blob/main/evaluate_enformer_pytorch_correlation.ipynb

A class to extract DNA sequences from FASTA files using pyfaidx. It handles intervals, chromosome boundaries, and padding with 'N's for sequences outside the defined range.

```python
class FastaStringExtractor:
    
    def __init__(self, fasta_file):
        self.fasta = pyfaidx.Fasta(fasta_file)
        self._chromosome_sizes = {k: len(v) for k, v in self.fasta.items()}

    def extract(self, interval: Interval, **kwargs) -> str:
        # Truncate interval if it extends beyond the chromosome lengths.
        chromosome_length = self._chromosome_sizes[interval.chrom]
        trimmed_interval = Interval(interval.chrom,
                                    max(interval.start, 0),
                                    min(interval.end, chromosome_length),
                                    )
        # pyfaidx wants a 1-based interval
        sequence = str(self.fasta.get_seq(trimmed_interval.chrom,
                                          trimmed_interval.start + 1,
                                          trimmed_interval.stop).seq).upper()
        # Fill truncated values with N's.
        pad_upstream = 'N' * max(-interval.start, 0)
        pad_downstream = 'N' * max(interval.end - chromosome_length, 0)
        return pad_upstream + sequence + pad_downstream

    def close(self):
        return self.fasta.close()
```

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