如何在Spacy Pipeline中移除段落/标题标记用Beacon令牌并保留Span属性

这问题我之前处理过类似的，spaCy的Doc对象是不可变的，所以不能直接删除令牌，但咱们可以通过创建新Doc的方式解决，同时保留你标记好的Span信息。下面给你一步步的解决方案：

第一步：确保信标被Tokenizer识别为单个令牌 #

默认情况下，spaCy的Tokenizer会把<p_start>拆成<、p_start、>三个令牌，这会导致后续匹配信标时出错。所以先给Tokenizer添加特殊规则，让信标被识别为单个令牌：

import spacy
from spacy.tokens import Doc, Span
from spacy.symbols import ORTH

# 加载预训练模型
nlp = spacy.load("en_core_web_sm")

# 为信标添加特殊分词规则
nlp.tokenizer.add_special_case("<p_start>", [{ORTH: "<p_start>"}])
nlp.tokenizer.add_special_case("<h_start>", [{ORTH: "<h_start>"}])

第二步：自定义组件标记段落/标题的Span #

这个组件会遍历分词后的令牌，找到信标位置，然后标记出对应的段落或标题范围：

def mark_section_spans(doc):
    # 初始化存储段落/标题的spans容器
    if "sections" not in doc.spans:
        doc.spans["sections"] = []
    
    i = 0
    while i < len(doc):
        token = doc[i]
        # 处理标题信标
        if token.text == "<h_start>":
            start = i + 1
            end = len(doc)
            # 找到下一个信标的位置作为当前标题的结束点
            for j in range(start, len(doc)):
                if doc[j].text in ("<h_start>", "<p_start>"):
                    end = j
                    break
            if start < end:
                span = Span(doc, start, end, label="HEADING")
                doc.spans["sections"].append(span)
            i = end
        # 处理段落信标
        elif token.text == "<p_start>":
            start = i + 1
            end = len(doc)
            for j in range(start, len(doc)):
                if doc[j].text in ("<h_start>", "<p_start>"):
                    end = j
                    break
            if start < end:
                span = Span(doc, start, end, label="PARAGRAPH")
                doc.spans["sections"].append(span)
            i = end
        else:
            i += 1
    return doc

第三步：自定义组件移除信标令牌并迁移Span信息 #

因为Doc是不可变对象，我们需要过滤掉信标令牌生成新的Doc，同时把之前标记的Span信息迁移过去（令牌索引变化后，Span的位置需要重新映射）：

def remove_beacon_tokens(doc):
    # 过滤掉所有信标令牌
    filtered_tokens = [t for t in doc if t.text not in ("<h_start>", "<p_start>")]
    
    # 用过滤后的令牌创建新Doc，保留原始的空格信息
    new_doc = Doc(
        doc.vocab,
        words=[t.text for t in filtered_tokens],
        spaces=[t.whitespace_ for t in filtered_tokens]
    )
    
    # 建立原令牌索引到新Doc索引的映射表
    idx_map = {}
    new_idx = 0
    for old_idx, token in enumerate(doc):
        if token.text not in ("<h_start>", "<p_start>"):
            idx_map[old_idx] = new_idx
            new_idx += 1
    
    # 迁移之前标记的Span到新Doc
    new_doc.spans["sections"] = []
    for span in doc.spans.get("sections", []):
        # 确保Span的起始和结束位置都在过滤后的令牌范围内
        if span.start in idx_map and (span.end - 1) in idx_map:
            new_start = idx_map[span.start]
            new_end = idx_map[span.end - 1] + 1  # Span的end是排他索引，需要+1
            new_span = Span(new_doc, new_start, new_end, label=span.label_)
            new_doc.spans["sections"].append(new_span)
    
    return new_doc

第四步：将组件添加到Pipeline的正确位置 #

把两个组件加到Tokenizer之后、Tok2Vec之前，保证执行顺序符合你的需求：

# 添加标记Span的组件
nlp.add_pipe("mark_section_spans", after="tokenizer")
# 添加移除信标的组件
nlp.add_pipe("remove_beacon_tokens", after="mark_section_spans")

测试验证 #

用示例文本测试效果：

test_text = "<h_start>机器学习入门指南<p_start>这是第一篇关于机器学习基础的介绍段落。<p_start>第二段会讲解常用的监督学习算法。"
doc = nlp(test_text)

print("处理后的纯文本:", doc.text)
print("\n识别到的段落/标题:")
for span in doc.spans["sections"]:
    print(f"{span.label_}: {span.text}")

输出会是不含信标的纯文本，同时Span也能正确标记对应的标题和段落内容。

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