清水泥沙

Transformer 学习之路 - 基于滑动窗口策略的机器阅读理解任务实现

Transformer 学习之路 - 基于滑动窗口策略的机器阅读理解任务实现

在自然语言处理（NLP）领域，机器阅读理解（Machine Reading Comprehension, MRC）是一个重要的任务，其目标是通过理解给定的上下文文本，回答用户提出的问题。本文将深入探讨如何利用 Transformer 模型，并结合滑动窗口策略来实现这一任务。

1. 背景与问题

机器阅读理解任务的核心挑战在于如何有效地处理长文本。传统的 Transformer 模型由于输入长度的限制，往往无法直接处理过长的上下文文本。为了解决这一问题，滑动窗口策略应运而生。通过将长文本分割成多个片段，并逐步滑动窗口进行处理，模型能够在不损失信息的情况下处理长文本。

2. 滑动窗口策略的实现

滑动窗口策略的核心思想是将长文本分割成多个重叠的片段，每个片段都包含一部分上下文信息。这样，模型可以在每个片段上独立地进行处理，最终将所有片段的结果进行整合。

2.1 数据预处理

首先，我们需要对数据进行预处理。以下代码展示了如何加载数据集，并使用滑动窗口策略对数据进行编码：

from datasets import load_dataset
from transformers import AutoTokenizer

# 加载数据集
datasets = load_dataset("cmrc2018", cache_dir="data")

# 加载预训练的 tokenizer
tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-macbert-base")

# 对数据进行编码，使用滑动窗口策略
tokenized_examples = tokenizer(
    text=sample_dataset["context"],
    text_pair=sample_dataset["question"],
    return_offsets_mapping=True,
    max_length=384,
    return_overflowing_tokens=True,  # 设置滑动窗口策略
    stride=128,  # 覆盖策略，每次移动128个字符
    truncation_strategy="only_second",
)

2.2 滑动窗口的处理

在处理滑动窗口时，我们需要确保每个片段都能够正确地定位答案的位置。以下代码展示了如何在滑动窗口的每个片段中定位答案的起始和结束位置：

for idx, _ in enumerate(sample_mapping):
    answer = sample_dataset["answers"][sample_mapping[idx]]
    start_char = answer["answer_start"][0]
    end_char = start_char + len(answer["text"][0])
    
    # 定位答案在token中的起始位置和结束位置
    context_start = tokenized_examples.sequence_ids(idx).index(1)
    context_end = tokenized_examples.sequence_ids(idx).index(None, context_start) - 1
    
    offset = tokenized_examples.get("offset_mapping")[idx]
    
    # 判断答案是否在context中
    if offset[context_end][1] < start_char or offset[context_start][0] > end_char:
        start_token_pos = 0
        end_token_pos = 0
    else:
        token_id = context_start
        while token_id <= context_end and offset[token_id][0] < start_char:
            token_id += 1
        start_token_pos = token_id
        token_id = context_end
        while token_id >= context_start and offset[token_id][1] > end_char:
            token_id -= 1
        end_token_pos = token_id
    
    print(answer, start_char, end_char, context_start, context_end, start_token_pos, end_token_pos)
    print("token answer decode:", tokenizer.decode(tokenized_examples["input_ids"][idx][start_token_pos: end_token_pos + 1]))

3. 模型训练与评估

在完成数据预处理后，我们可以开始训练模型。以下代码展示了如何加载模型、配置训练参数，并进行训练：

from transformers import AutoModelForQuestionAnswering, Trainer, TrainingArguments

# 加载预训练模型
model = AutoModelForQuestionAnswering.from_pretrained("hfl/chinese-macbert-base")

# 配置训练参数
args = TrainingArguments(
    output_dir="models_for_qa",
    per_device_train_batch_size=50,
    per_device_eval_batch_size=50,
    eval_strategy="steps",
    eval_steps=200,
    save_strategy="epoch",
    logging_steps=50,
    num_train_epochs=1
)

# 创建 Trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=DefaultDataCollator(),
    tokenizer=tokenizer,
    compute_metrics=metirc
)

# 训练模型
trainer.train()

3.1 模型评估

在训练完成后，我们需要对模型进行评估。以下代码展示了如何获取模型的输出，并进行评估：

import numpy as np
import collections

def get_result(start_logits, end_logits, examples, features):
    predictions = {}
    references = {}
    
    # example 和 feature的映射
    example_to_feature = collections.defaultdict(list)
    for idx, example_id in enumerate(features["example_ids"]):
        example_to_feature[example_id].append(idx)
    
    # 最优答案候选
    n_best = 20
    # 最大答案长度
    max_answer_length = 30
    
    for example in examples:
        example_id = example["id"]
        context = example["context"]
        answers = []
        for feature_idx in example_to_feature[example_id]:
            start_logit = start_logits[feature_idx]
            end_logit = end_logits[feature_idx]
            offset = features[feature_idx]["offset_mapping"]
            start_indexes = np.argsort(start_logit)[::-1][:n_best].tolist()
            end_indexes = np.argsort(end_logit)[::-1][:n_best].tolist()
            for start_index in start_indexes:
                for end_index in end_indexes:
                    if offset[start_index] is None or offset[end_index] is None:
                        continue
                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:
                        continue
                    answers.append({
                        "text": context[offset[start_index][0]: offset[end_index][1]],
                        "score": start_logit[start_index] + end_logit[end_index]
                    })
        if len(answers) > 0:
            best_answer = max(answers, key=lambda x: x["score"])
            predictions[example_id] = best_answer["text"]
        else:
            predictions[example_id] = ""
        references[example_id] = example["answers"]["text"]
    
    return predictions, references

4. 模型预测与保存

在模型训练和评估完成后，我们可以使用模型进行预测，并将模型保存以备后续使用：

from transformers import pipeline

# 使用模型进行预测
qa = pipeline("question-answering", model=model, tokenizer=tokenizer)
result = qa(question="信宜在哪里", context="广东茂名信宜的特色是三华李炒猪大肠。")
print(result)

# 保存模型
model_save_path = "/content/drive/MyDrive/ai-learning/2.NLP Task/03-question_answering/mrc2"
model.save_pretrained(model_save_path)

5. 总结

通过本文的讲解，我们深入探讨了如何利用 Transformer 模型和滑动窗口策略来实现机器阅读理解任务。滑动窗口策略有效地解决了长文本处理的问题，使得模型能够在有限的输入长度内处理更长的上下文。结合代码示例，我们展示了从数据预处理、模型训练到模型预测的完整流程，希望能够帮助读者更好地理解和应用 Transformer 技术。