HuggingFace：使用 Transformer 对 DNA 序列进行高效大规模嵌入提取

我有一个非常大的数据框（60+ 百万行），我想使用转换器模型来获取这些行（DNA 序列）的嵌入。基本上，这首先涉及标记化，然后我可以获得嵌入。 由于 RAM 限制，我发现标记化然后将所有内容嵌入到一个 py 文件中是行不通的。这是我发现的解决方法，适用于大约 3000 万行的数据帧（但不适用于较大的 df）：

1. 标记化 - 将输出保存为 200 个块/分片
2. 将这 200 个块分别提供给获取嵌入
3. 这些嵌入，然后连接成一个更大的嵌入文件

最终嵌入文件应包含以下列： [['Cromosome', 'label', 'embeddings']]

总的来说，我对如何让它适用于我的更大数据集有点迷失。

我已经研究过流媒体数据集，但我认为这实际上没有帮助，因为我需要所有嵌入，而不仅仅是一些嵌入。如果我流式传输标记化并将其一次一点地输入到嵌入过程中（沿途删除标记），也许它可以工作。这样，我就不必保存令牌了。如果这不可行，请纠正我。

理想情况下，我希望避免对数据进行分片，但我只是希望代码在此时运行而不达到 RAM 限制。

步骤 1|| |第 2 步（此代码在 200 个分片中运行）

dataset = Dataset.from_pandas(element_final[['Chromosome', 'sequence', 'label']]) 

dataset = dataset.shuffle(seed=42)
tokenizer = AutoTokenizer.from_pretrained(f"InstaDeepAI/nucleotide-transformer-500m-human-ref")
def tokenize_function(examples):
    outputs = tokenizer.batch_encode_plus(examples["sequence"], return_tensors="pt", truncation=False, padding=False, max_length=80)
    return outputs
    
# Creating tokenized  dataset
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True, batch_size=2000)

tokenized_dataset.save_to_disk(f"tokenized_elements/tokenized_{ELEMENT}", num_shards=200)

第 3 步：将所有 200 个嵌入分片连接成 1 个分片。

input_file = f"tokenized_elements/tokenized_{ELEMENT_LABEL}/{filename}.arrow"

# Load input data
d1 = Dataset.from_file(input_file)

def embed_function(examples):
    torch.cuda.empty_cache()
    gc.collect()

    inputs = torch.tensor(examples['input_ids'])  # Convert to tensor
    inputs = inputs.to(device)

    with torch.no_grad():
        outputs = model(input_ids=inputs, output_hidden_states=True)

    # Step 3: Extract the embeddings
    hidden_states = outputs.hidden_states  # List of hidden states from all layers
    embeddings = hidden_states[-1]  # Assuming you want embeddings from the last layer
    averaged_embeddings = torch.mean(embeddings, dim=1)  # Calculate mean along dimension 1 (the dimension with size 86)
    averaged_embeddings = averaged_embeddings.to(torch.float32)  # Ensure float32 data type
    return {'embeddings': averaged_embeddings}

# Map embeddings function to input data
embeddings = d1.map(embed_function, batched=True, batch_size=1550)
embeddings = embeddings.remove_columns(["input_ids", "attention_mask"])

# Save embeddings to disk
output_dir = f"embedded_elements/embeddings_{ELEMENT_LABEL}/{filename}"  # Assuming ELEMENT_LABEL is defined elsewhere

step 3: concatenate all 200 shards of embeddings into 1.

以下是如何使用 Hugging Face Transformers 有效地从 DNA 序列中提取大规模嵌入的方法，并解决遇到的内存问题：

import pandas as pd
from transformers import AutoTokenizer, AutoModel
import torch
from torch.utils.data import Dataset, DataLoader
import gc

# 定义设备（CPU 或 GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义模型和分词器
model_name = "InstaDeepAI/nucleotide-transformer-500m-human-ref" 
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name).to(device)

# 定义一个自定义数据集类
class DNADataset(Dataset):
    def __init__(self, dataframe, tokenizer, max_length):
        self.sequences = dataframe['sequence'].tolist()
        self.labels = dataframe['label'].tolist()
        self.chromosomes = dataframe['Chromosome'].tolist()
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        sequence = self.sequences[idx]
        label = self.labels[idx]
        chromosome = self.chromosomes[idx]

        encoding = self.tokenizer(
            sequence,
            add_special_tokens=True,
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt",
        )

        return {
            'input_ids': encoding['input_ids'].squeeze(0),
            'attention_mask': encoding['attention_mask'].squeeze(0),
            'label': torch.tensor(label),
            'chromosome': chromosome
        }

# 定义数据加载器和嵌入提取函数
def get_embeddings(model, dataloader, device):
    model.eval()
    all_embeddings = []
    all_labels = []
    all_chromosomes = []
    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            embeddings = outputs.last_hidden_state[:, 0, :].cpu() # 使用 [CLS] 标记的嵌入

            all_embeddings.append(embeddings)
            all_labels.extend(batch['label'])
            all_chromosomes.extend(batch['chromosome'])

    return torch.cat(all_embeddings), all_labels, all_chromosomes

# 加载数据帧（分块加载以节省内存）
chunksize = 1000000  # 根据需要调整分块大小
embeddings_list = []
labels_list = []
chromosomes_list = []

for chunk in pd.read_csv("your_data.csv", chunksize=chunksize):
    # 创建数据集和数据加载器
    dataset = DNADataset(chunk, tokenizer, max_length=80)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=False) # 根据需要调整批大小

    # 提取嵌入
    embeddings, labels, chromosomes = get_embeddings(model, dataloader, device)

    # 将结果存储到列表中
    embeddings_list.append(embeddings)
    labels_list.extend(labels)
    chromosomes_list.extend(chromosomes)

    # 释放内存
    del dataset, dataloader, embeddings, labels, chromosomes
    gc.collect()
    torch.cuda.empty_cache()

# 连接所有分块的结果
all_embeddings = torch.cat(embeddings_list)
all_labels = labels_list
all_chromosomes = chromosomes_list

# 创建最终数据帧
final_df = pd.DataFrame({
    'Chromosome': all_chromosomes,
    'label': all_labels,
    'embeddings': all_embeddings.tolist() # 将张量转换为列表以存储到数据帧
})

# 保存最终数据帧
final_df.to_csv("embeddings.csv", index=False)

说明：

1. 数据分块： 该代码使用 pandas.read_csv 的 chunksize 参数将大型数据帧加载到内存中。这允许以较小的块处理数据，从而防止内存溢出。
2. 自定义数据集和数据加载器： 使用 torch.utils.data.Dataset 和 torch.utils.data.DataLoader 可以有效地加载和预处理数据。这使可以对批处理、混洗和将数据加载到设备（CPU/GPU）进行细粒度控制。
3. GPU 加速（如果可用）： 代码会自动检测并使用可用的 GPU 来加速嵌入提取过程。
4. 批处理： 嵌入提取在批处理模式下执行，这有助于提高效率并减少内存消耗。
5. 内存管理： 该代码使用 del 、 gc.collect() 和 torch.cuda.empty_cache() 在每个分块处理后显式释放内存。这有助于最大限度地减少内存使用并防止内存泄漏。

其他提示：

• 调整批大小和分块大小： 根据的硬件资源和数据集大小，可能需要调整 batch_size 和 chunksize 参数以优化性能和内存使用。
• 使用更小的模型： 如果仍然遇到内存问题，请考虑使用更小、更快的 Transformer 模型。
• 降维： 可以使用 PCA 或 t-SNE 等降维技术来减少嵌入的维数。这可以帮助减少内存使用，而不会牺牲太多信息。

通过遵循这些步骤，应该能够有效地从大型 DNA 序列数据集中提取嵌入，而不会遇到内存问题。