BERT词嵌入长文本处理与内存优化实践

import torch
from transformers import AutoModel, AutoTokenizer

# 示例文本数据
texts = [
    "这是一个示例文本，用于演示如何生成BERT词嵌入。",
    "另一个更长的文本，需要进行截断以适应模型的最大序列长度限制，同时确保内存不会溢出。",
    # ... 更多文本数据
]

# 加载匹配的模型和分词器
# 可以根据需求选择不同的预训练模型，例如 "bert-base-uncased"
model_name = "indolem/indobert-base-uncased" # 示例模型
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 如果有GPU可用，将模型移动到GPU
if torch.cuda.is_available():
    model.to('cuda')
    print("模型已加载到GPU。")
else:
    print("GPU不可用，模型将在CPU上运行。")

# 对批量的句子进行分词，设置最大序列长度并进行截断和填充
# max_length: 模型的最大输入序列长度，例如512
# truncation=True: 当文本长度超过max_length时，自动截断
# padding=True: 将所有序列填充到批次中最长序列的长度（或max_length，如果max_length更短）
# return_tensors='pt': 返回PyTorch张量
tokenized_texts = tokenizer(texts,
                            max_length=512,
                            truncation=True,
                            padding=True,
                            return_tensors='pt')

print(f"分词后的输入ID形状: {tokenized_texts['input_ids'].shape}")
print(f"分词后的注意力掩码形状: {tokenized_texts['attention_mask'].shape}")

# 将分词结果移动到GPU（如果模型在GPU上）
if torch.cuda.is_available():
    input_ids = tokenized_texts['input_ids'].to('cuda')
    attention_mask = tokenized_texts['attention_mask'].to('cuda')
else:
    input_ids = tokenized_texts['input_ids']
    attention_mask = tokenized_texts['attention_mask']

# 前向传播获取词嵌入
with torch.no_grad():
    outputs = model(input_ids=input_ids,
                    attention_mask=attention_mask)

    # 提取词嵌入，通常是模型的最后一层隐藏状态
    # 形状为 [batch_size, num_seq_tokens, embed_size]
    word_embeddings = outputs.last_hidden_state

print(f"生成的词嵌入形状: {word_embeddings.shape}")
# 示例输出: torch.Size([2, 512, 768])，表示2个样本，每个样本512个token，每个token有768维的嵌入

# 假设有一个很长的文本列表 all_texts
all_texts = ['text1', 'text2', ..., 'textN']
batch_size = 8 # 根据GPU显存大小调整，可以尝试更小的值如4, 2, 1

all_word_embeddings = []

for i in range(0, len(all_texts), batch_size):
    current_batch_texts = all_texts[i : i + batch_size]

    tokenized_batch = tokenizer(current_batch_texts,
                                max_length=512,
                                truncation=True,
                                padding=True,
                                return_tensors='pt')

    if torch.cuda.is_available():
        input_ids_batch = tokenized_batch['input_ids'].to('cuda')
        attention_mask_batch = tokenized_batch['attention_mask'].to('cuda')
    else:
        input_ids_batch = tokenized_batch['input_ids']
        attention_mask_batch = tokenized_batch['attention_mask']

    with torch.no_grad():
        outputs_batch = model(input_ids=input_ids_batch,
                              attention_mask=attention_mask_batch)
        word_embeddings_batch = outputs_batch.last_hidden_state
        all_word_embeddings.append(word_embeddings_batch.cpu()) # 将结果移回CPU以释放GPU内存

# 如果需要，可以将所有批次的词嵌入拼接起来
# final_embeddings = torch.cat(all_word_embeddings, dim=0)
# print(f"所有文本的最终词嵌入形状: {final_embeddings.shape}")