使用HuggingFace预训练模型生成文档嵌入时的张量尺寸错误

问题 #

使用HuggingFace Transformers库的distilbert-base-cased模型提取PDF文档嵌入时，因文本长度超出模型最大序列限制报错，具体表现为序列长度3619超过模型允许的512长度，触发张量维度不匹配的RuntimeError。

代码 #

from transformers import pipeline, AutoTokenizer, AutoModel
from transformers import RobertaTokenizer, RobertaModel
import fitz
from openpyxl import load_workbook
import os
from tqdm import tqdm

PRETRAIN_MODEL = 'distilbert-base-cased'
DIR = "dataset"

# Load and process the text
all_files = os.listdir(DIR)
pdf_texts = {}
for filename in all_files:
    if filename.lower().endswith('.pdf'):
        pdf_path = os.path.join(DIR, filename)
        with fitz.open(pdf_path) as doc:
            text_content = ""
            for page in doc:
                text_content += page.get_text()
            text = text_content.split("PUBLIC CONSULTATION")[0]
            project_code = os.path.splitext(filename)[0]
            pdf_texts[project_code] = text 

# Generate embeddings for the documents
tokenizer = AutoTokenizer.from_pretrained(PRETRAIN_MODEL)
model = AutoModel.from_pretrained(PRETRAIN_MODEL)
pipe = pipeline('feature-extraction', model=model, tokenizer=tokenizer)

embeddings = {}
for project_code, text in tqdm(pdf_texts.items(), desc="Generating embeddings", unit="doc"):
    embedding = pipe(text, return_tensors="pt")
    embeddings[project_code] = embedding[0][0].numpy()

错误信息 #

Generating embeddings:   0%|          | 0/58 [00:00<?, ?doc/s]Token indices sequence length is longer than the specified maximum sequence length for this model (3619 > 512). Running this sequence through the model will result in indexing errors
Generating embeddings:   0%|          | 0/58 [00:00<?, ?doc/s]
RuntimeError: The size of tensor a (3619) must match the size of tensor b (512) at non-singleton dimension 1

解决方案 #

方案1：截断过长文本 #

直接对超出模型最大长度的文本进行截断，保留前512个token，修改pipeline调用参数即可：

embeddings = {}
for project_code, text in tqdm(pdf_texts.items(), desc="Generating embeddings", unit="doc"):
    # 添加截断参数，强制文本长度不超过模型限制
    embedding = pipe(text, return_tensors="pt", truncation=True, max_length=512)
    embeddings[project_code] = embedding[0][0].numpy()

方案2：分段处理长文本并聚合嵌入 #

若不想丢失文本信息，可将长文本分割为多个512token的片段，分别生成嵌入后取平均值作为文档整体嵌入：

import torch

def get_document_embedding(text, model, tokenizer, max_length=512):
    # 分割文本为多个符合长度要求的片段
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_length, padding=False, return_overflowing_tokens=True)
    del inputs["overflow_to_sample_mapping"]  # 移除多余字段避免模型报错
    with torch.no_grad():
        outputs = model(**inputs)
    # 取每个片段<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token的嵌入，再计算平均值
    cls_embeddings = outputs.last_hidden_state[:, 0, :]
    doc_embedding = torch.mean(cls_embeddings, dim=0).numpy()
    return doc_embedding

# 替换原嵌入生成逻辑
embeddings = {}
for project_code, text in tqdm(pdf_texts.items(), desc="Generating embeddings", unit="doc"):
    embedding = get_document_embedding(text, model, tokenizer)
    embeddings[project_code] = embedding

方案3：使用支持更长序列的模型 #

换用原生支持超长序列的预训练模型，比如allenai/longformer-base-4096（支持4096个token），仅需修改模型名称：

PRETRAIN_MODEL = 'allenai/longformer-base-4096'

内容的提问来源于stack exchange，提问作者Howie