如何使用Hugging Face 512最大token的NER模型处理长文本？

核心原理 #

BERT系列模型的最大输入token数限制为512，超过该长度会触发维度不匹配报错。解决思路是将长文本拆分为多个不超过512token的子块，同时保留重叠区域避免实体被截断，最后合并各子块的识别结果。

具体实现步骤 #

1. 编写文本拆分函数 #

实现一个函数，按token拆分长文本，每个子块预留[CLS]和[SEP]的位置（各占1个token），同时设置重叠token段，确保跨块实体能被完整识别：

def split_long_text(text, tokenizer, max_length=512, overlap=100):
    # 编码文本得到不含特殊token的token列表
    tokens = tokenizer.encode(text, add_special_tokens=False)
    chunks = []
    start_idx = 0
    total_tokens = len(tokens)

    while start_idx < total_tokens:
        # 计算当前块的结束位置，预留2个token给[CLS]和[SEP]
        end_idx = start_idx + (max_length - 2)
        if end_idx > total_tokens:
            end_idx = total_tokens
        # 将token转换回文本块
        chunk_tokens = tokens[start_idx:end_idx]
        chunk_text = tokenizer.decode(chunk_tokens)
        chunks.append(chunk_text)
        # 更新起始位置，减去重叠长度避免实体截断
        start_idx = end_idx - overlap
        # 防止极端情况导致死循环
        if start_idx >= total_tokens:
            break
    return chunks

2. 批量处理拆分后的文本块 #

初始化模型和NER pipeline，遍历所有子块并收集识别结果：

model_name = "cahya/bert-base-indonesian-NER"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForTokenClassification.from_pretrained(model_name)
nlp = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")

# 替换为你的超长印尼语文本
long_text = "此处输入你的超长印尼语文本内容..."

# 拆分长文本
text_chunks = split_long_text(long_text, tokenizer)

# 批量处理并收集所有实体
all_entities = []
for chunk in text_chunks:
    chunk_entities = nlp(chunk)
    all_entities.extend(chunk_entities)

3. 去重合并结果（可选） #

由于重叠区域可能导致同一实体被重复识别，可按实体文本和类型去重：

unique_entities = []
seen_entities = set()
for ent in all_entities:
    # 用实体文本+类型作为去重键
    entity_key = (ent["word"], ent["entity_group"])
    if entity_key not in seen_entities:
        seen_entities.add(entity_key)
        unique_entities.append(ent)

关键注意事项 #

• 重叠长度调整：重叠token数可根据文本中实体的平均长度调整，建议设置为50-200，平衡识别完整性和效率。
• 位置映射（可选）：如果需要保留实体在原文本中的精确位置，拆分时需记录每个子块在原文本的起始偏移量，最后通过token与原文本的映射关系还原位置。

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