问答系统：阅读理解与开放域 QA

# QA 任务分类与形式化定义
from dataclasses import dataclass
from typing import List, Optional

@dataclass
class QAProblem:
    """QA 问题的统一表示"""
    question: str            # 问题
    context: Optional[str]   # 上下文（抽取式/生成式有，开放域无）
    answer: Optional[str]    # 标准答案
    answer_start: Optional[int]  # 抽取式的起始位置
    answer_end: Optional[int]    # 抽取式的结束位置

class ExtractiveQA:
    """抽取式 QA：从上下文中定位答案子串"""
    
    def predict(self, question: str, context: str) -> dict:
        """
        返回: {"answer": str, "start": int, "end": int, "score": float}
        模型预测答案在 context 中的起始和结束位置
        """
        # BERT 输出 start_logits 和 end_logits
        # answer = context[start:end]
        pass

class GenerativeQA:
    """生成式 QA：自由生成答案文本"""
    
    def predict(self, question: str, context: Optional[str] = None) -> str:
        """
        模型可以基于 context 生成答案，也可以仅基于问题生成
        适用于 T5/BART 等 Seq2Seq 模型
        """
        pass

class OpenDomainQA:
    """开放域 QA：检索+阅读两阶段"""
    
    def predict(self, question: str) -> str:
        """
        1. Retriever: 从大规模语料中检索 top-K 文档
        2. Reader: 从 top-K 文档中提取/生成答案
        """
        pass

# 任务分类对比
tasks = {
    "抽取式": {"上下文": "给定", "答案来源": "上下文子串", "评估": "精确匹配/ F1"},
    "生成式": {"上下文": "可选", "答案来源": "自由生成", "评估": "人工评估/ BLEU"},
    "开放域": {"上下文": "无", "答案来源": "检索+提取", "评估": "EM/F1/召回率"},
}
for task, info in tasks.items():
    print(f"{task}: {info}")

# 问答系统评估指标详解
import re
import string
from collections import Counter

def normalize_answer(text: str) -> str:
    """标准化答案文本：转小写、去标点、去多余空白"""
    def remove_articles(t):
        return re.sub(r'\b(a|an|the)\b', ' ', t)
    def white_space_fix(t):
        return ' '.join(t.split())
    def remove_punc(t):
        exclude = set(string.punctuation)
        return ''.join(ch for ch in t if ch not in exclude)
    def lower(t):
        return t.lower()
    return white_space_fix(remove_articles(remove_punc(lower(text))))

def exact_match(prediction: str, ground_truth: str) -> float:
    """精确匹配（EM）：预测答案与标准答案完全一致"""
    return float(normalize_answer(prediction) == normalize_answer(ground_truth))

def f1_score(prediction: str, ground_truth: str) -> float:
    """F1 分数：基于 token 级别的重叠度"""
    pred_tokens = normalize_answer(prediction).split()
    gt_tokens = normalize_answer(ground_truth).split()
    
    if len(pred_tokens) == 0 or len(gt_tokens) == 0:
        return float(pred_tokens == gt_tokens)
    
    common = Counter(pred_tokens) & Counter(gt_tokens)
    num_same = sum(common.values())
    
    if num_same == 0:
        return 0.0
    
    precision = num_same / len(pred_tokens)
    recall = num_same / len(gt_tokens)
    return 2 * (precision * recall) / (precision + recall)

# SQuAD 评估示例
print("=== SQuAD 评估示例 ===")
predictions = [
    ("the United States", "United States"),
    ("in 1990", "1990"),
    ("a neural network", "neural networks"),
    ("Paris", "London"),
]

for pred, gold in predictions:
    em = exact_match(pred, gold)
    f1 = f1_score(pred, gold)
    print(f"预测: "{pred}" | 标准: "{gold}" | EM: {em:.0f} | F1: {f1:.2f}")

# SQuAD 数据集加载与分析
from datasets import load_dataset

# 加载 SQuAD v2（包含无解问题）
squad = load_dataset("squad_v2")

print("=== SQuAD 2.0 数据集统计 ===")
print(f"训练集样本数: {len(squad['train'])}")
print(f"验证集样本数: {len(squad['validation'])}")

# 分析数据分布
train_data = squad['train']
answerable_count = sum(1 for a in train_data['answers'] if a['text'])
unanswerable_count = len(train_data) - answerable_count
print(f"\n可回答问题数: {answerable_count:,}")
print(f"无解问题数: {unanswerable_count:,}")
print(f"无解问题比例: {unanswerable_count/len(train_data)*100:.1f}%")

# 答案长度分布
answer_lengths = [len(a['text'][0]) if a['text'] else 0 for a in train_data['answers']]
avg_length = sum(l for l in answer_lengths if l > 0) / max(answerable_count, 1)
print(f"\n平均答案长度: {avg_length:.1f} 字符")
print(f"最短答案: {min(l for l in answer_lengths if l > 0)} 字符")
print(f"最长答案: {max(answer_lengths)} 字符")

# 查看一个样本
sample = train_data[0]
print(f"\n=== 示例 ===")
print(f"标题: {sample['title']}")
print(f"上下文: {sample['context'][:200]}...")
print(f"问题: {sample['question']}")
print(f"答案: {sample['answers']['text'][:3]}")

# SQuAD 格式的 JSON 结构解析
import json

# SQuAD 原始 JSON 格式
squad_json = {
    "version": "2.0",
    "data": [
        {
            "title": "Normans",
            "paragraphs": [
                {
                    "context": "The Normans (Norman: Nourmands; French: Normands; Latin: Normanni) were the people who in the 10th and 11th centuries gave their name to Normandy, a region in France.",
                    "qas": [
                        {
                            "question": "In what country is Normandy located?",
                            "id": "56ddde6b9a695914005b9628",
                            "answers": [
                                {
                                    "text": "France",
                                    "answer_start": 159
                                }
                            ],
                            "is_impossible": false
                        },
                        {
                            "question": "When were the Normans in Normandy?",
                            "id": "56ddde6b9a695914005b9629",
                            "answers": [],
                            "is_impossible": true
                        }
                    ]
                }
            ]
        }
    ]
}

# 展平为 (context, question, answer) 三元组
def flatten_squad(data):
    """将嵌套的 SQuAD JSON 展平为扁平列表"""
    records = []
    for article in data["data"]:
        title = article["title"]
        for paragraph in article["paragraphs"]:
            context = paragraph["context"]
            for qa in paragraph["qas"]:
                is_impossible = qa.get("is_impossible", False)
                if not is_impossible and qa["answers"]:
                    answer = qa["answers"][0]
                    records.append({
                        "title": title,
                        "context": context,
                        "question": qa["question"],
                        "answer_text": answer["text"],
                        "answer_start": answer["answer_start"],
                    })
    return records

records = flatten_squad(squad_json)
for r in records:
    print(f"Q: {r['question']}")
    print(f"A: "{r['answer_text']}" (start={r['answer_start']})")
    print(f"Context preview: ...{r['context'][max(0,r['answer_start']-20):r['answer_start']+len(r['answer_text'])+20]}...")
    print()

# BERT 抽取式 QA 模型实现
import torch
import torch.nn as nn
from transformers import BertModel

class BertForExtractiveQA(nn.Module):
    """基于 BERT 的抽取式问答模型"""
    
    def __init__(self, model_name="bert-base-chinese"):
        super().__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.qa_start = nn.Linear(self.bert.config.hidden_size, 1)
        self.qa_end = nn.Linear(self.bert.config.hidden_size, 1)
    
    def forward(self, input_ids, attention_mask,
                start_positions=None, end_positions=None):
        """
        input_ids: (batch, seq_len)
        attention_mask: (batch, seq_len)
        start_positions: (batch,) 答案起始位置
        end_positions: (batch,) 答案结束位置
        """
        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            return_dict=True
        )
        sequence_output = outputs.last_hidden_state  # (batch, seq_len, hidden)
        
        # 预测起始和结束位置的 logits
        start_logits = self.qa_start(sequence_output).squeeze(-1)  # (batch, seq_len)
        end_logits = self.qa_end(sequence_output).squeeze(-1)      # (batch, seq_len)
        
        total_loss = None
        if start_positions is not None and end_positions is not None:
            # 如果位置超出序列长度，截断到序列末尾
            if start_positions.max() >= start_logits.shape[1]:
                start_positions = torch.clamp(
                    start_positions, max=start_logits.shape[1] - 1
                )
            if end_positions.max() >= end_logits.shape[1]:
                end_positions = torch.clamp(
                    end_positions, max=end_logits.shape[1] - 1
                )
            
            start_loss = nn.CrossEntropyLoss()(start_logits, start_positions)
            end_loss = nn.CrossEntropyLoss()(end_logits, end_positions)
            total_loss = (start_loss + end_loss) / 2
        
        return {
            "loss": total_loss,
            "start_logits": start_logits,
            "end_logits": end_logits,
        }

    def extract_answer(self, input_ids, attention_mask, tokenizer):
        """推理：提取答案"""
        outputs = self(input_ids, attention_mask)
        start_logits = outputs["start_logits"]
        end_logits = outputs["end_logits"]
        
        # 获取最高概率的起止位置
        start_idx = torch.argmax(start_logits, dim=-1).item()
        end_idx = torch.argmax(end_logits, dim=-1).item()
        
        # 确保 end >= start
        if end_idx < start_idx:
            end_idx = start_idx
        
        answer_ids = input_ids[0][start_idx:end_idx + 1]
        answer = tokenizer.decode(answer_ids, skip_special_tokens=True)
        return answer, start_idx, end_idx

# BERT QA 推理流程演示
from transformers import BertTokenizer
import torch

# 加载模型和分词器
model_name = "bert-base-uncased"
tokenizer = BertTokenizer.from_pretrained(model_name)

# 示例问题与上下文
question = "What year did the Normans give their name to Normandy?"
context = "The Normans (Norman: Nourmands; French: Normands; Latin: Normanni) were the people who in the 10th and 11th centuries gave their name to Normandy, a region in France."

# 编码输入
inputs = tokenizer(
    question, context,
    return_tensors="pt",
    max_length=512,
    truncation=True,
    padding="max_length"
)

print("=== BERT QA 推理 ===")
print(f"问题: {question}")
print(f"上下文: {context}")
print(f"\n输入序列长度: {inputs['input_ids'].shape[1]} tokens")

# 解码 token 序列以可视化
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
sep_indices = [i for i, t in enumerate(tokens) if t == '[SEP]']
print(f"\nToken 序列可视化:")
print(f"  [CLS] + 问题 token → [SEP] (索引 {sep_indices[0]}) + 上下文 token → [SEP]")

# 模拟 BERT 输出（实际应调用模型）
# start_logits 和 end_logits 由 BERT + QA heads 生成
print(f"\n=== 跨度预测 ===")
print("BERT 对每个 token 输出两个分数:")
print("  start_score[token_i] = W_start · H_i")
print("  end_score[token_i] = W_end · H_i")
print("答案 = tokens[argmax(start_score) : argmax(end_score) + 1]")

# 长上下文窗口切分策略
def split_context_into_windows(context, tokenizer, max_len=512, stride=128):
    """将长上下文切分为重叠窗口"""
    tokens = tokenizer.encode(context, add_special_tokens=False)
    windows = []
    for i in range(0, len(tokens), max_len - stride):
        window = tokens[i:i + max_len]
        windows.append({
            "tokens": window,
            "offset": i,
            "length": len(window)
        })
    return windows

# 模拟窗口切分
windows = split_context_into_windows("A" * 2000, tokenizer, max_len=512, stride=128)
print(f"\n=== 长上下文处理 ===")
print(f"原文长度: 2000 字符")
print(f"窗口数: {len(windows)}")
print(f"窗口大小: 512 tokens, 步长: 128 tokens")

# DPR (Dense Passage Retrieval) 检索器实现
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BertModel, BertTokenizer

class DPRRetriever(nn.Module):
    """密集段落检索器"""
    
    def __init__(self, model_name="bert-base-uncased", projection_dim=768):
        super().__init__()
        self.question_encoder = BertModel.from_pretrained(model_name)
        self.context_encoder = BertModel.from_pretrained(model_name)
        self.projection_q = nn.Linear(768, projection_dim)
        self.projection_c = nn.Linear(768, projection_dim)
    
    def encode_question(self, input_ids, attention_mask):
        """将问题编码为稠密向量"""
        outputs = self.question_encoder(input_ids, attention_mask=attention_mask)
        pooled = outputs.pooler_output  # (batch, 768)
        return self.projection_q(pooled)  # (batch, projection_dim)
    
    def encode_context(self, input_ids, attention_mask):
        """将文档编码为稠密向量"""
        outputs = self.context_encoder(input_ids, attention_mask=attention_mask)
        pooled = outputs.pooler_output
        return self.projection_c(pooled)
    
    def contrastive_loss(self, q_embeds, pos_c_embeds, neg_c_embeds, temperature=0.07):
        """
        对比损失：拉近问题-正文档对，推远问题-负文档对
        q_embeds: (batch, d) 问题向量
        pos_c_embeds: (batch, d) 正文档向量
        neg_c_embeds: (batch, K, d) 负文档向量
        """
        # 计算正样本相似度
        pos_sim = torch.sum(q_embeds * pos_c_embeds, dim=-1) / temperature  # (batch,)
        
        # 计算负样本相似度
        neg_sim = torch.einsum('bd,bkd->bk', q_embeds, neg_c_embeds) / temperature
        
        # 拼接 logits
        logits = torch.cat([pos_sim.unsqueeze(1), neg_sim], dim=-1)  # (batch, K+1)
        labels = torch.zeros(q_embeds.shape[0], dtype=torch.long, device=q_embeds.device)
        
        return F.cross_entropy(logits, labels)

# FAISS 快速相似度检索
print("=== FAISS 检索流程 ===")
print("1. 将所有文档编码为稠密向量并构建 FAISS 索引")
print("2. 将问题编码为稠密向量")
print("3. 在 FAISS 索引中进行最近邻搜索")
print("4. 返回 top-K 相关文档")

# 开放域 QA 完整两阶段流程
import torch
import numpy as np

class OpenDomainQA:
    """开放域问答系统：检索 + 阅读"""
    
    def __init__(self, retriever, reader, document_index):
        """
        retriever: DPRRetriever 实例
        reader: BertForExtractiveQA 或 T5 实例
        document_index: 文档集合 + FAISS 索引
        """
        self.retriever = retriever
        self.reader = reader
        self.doc_index = document_index
        self.top_k = 5  # 检索 top-5 文档
    
    def answer(self, question: str) -> dict:
        """
        1. 编码问题 → 稠密向量
        2. FAISS 检索 → top-K 文档
        3. 阅读器对每个文档做阅读理解
        4. 聚合多个文档的答案
        """
        # 阶段 1: 检索
        q_embed = self.retriever.encode_question(question)
        doc_ids, scores = self.doc_index.search(q_embed, k=self.top_k)
        
        # 阶段 2: 阅读
        candidates = []
        for doc_id, score in zip(doc_ids, scores):
            context = self.doc_index.get_document(doc_id)
            answer, confidence = self.reader.predict(question, context)
            candidates.append({
                "doc_id": doc_id,
                "doc_title": self.doc_index.get_title(doc_id),
                "answer": answer,
                "retrieval_score": float(score),
                "reading_confidence": float(confidence),
                "combined_score": float(score) * float(confidence),
            })
        
        # 按综合分数排序
        candidates.sort(key=lambda x: x["combined_score"], reverse=True)
        return {
            "question": question,
            "best_answer": candidates[0]["answer"],
            "candidates": candidates,
        }

# 模拟执行
print("=== 开放域 QA 执行示例 ===")
question = "Who painted the Mona Lisa?"
print(f"问题: {question}")
print(f"\n阶段 1 - 检索:")
print(f"  编码问题 → 768 维向量")
print(f"  FAISS 搜索 → 检索 top-5 文档")
print(f"\n阶段 2 - 阅读:")
print(f"  文档1: Wikipedia - Leonardo da Vinci → 答案: Leonardo da Vinci")
print(f"  文档2: Wikipedia - Mona Lisa → 答案: Leonardo da Vinci")
print(f"  文档3: Wikipedia - Italian Renaissance → 答案: 无")
print(f"  文档4: Wikipedia - Louvre Museum → 答案: 无")
print(f"  文档5: Wikipedia - Renaissance Art → 答案: 无")
print(f"\n最终答案: Leonardo da Vinci")
print(f"置信度: 0.97 (多文档一致)")

# RAG (Retrieval-Augmented Generation) 实现
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
import torch

# 加载预训练 RAG 模型
model_name = "facebook/rag-sequence-nq"  # NQ 变体
tokenizer = RagTokenizer.from_pretrained(model_name)
retriever = RagRetriever.from_pretrained(model_name)
model = RagSequenceForGeneration.from_pretrained(model_name, retriever=retriever)

# RAG 问答
question = "What is the capital of France?"
input_ids = tokenizer(question, return_tensors="pt").input_ids

# 生成答案
outputs = model.generate(input_ids=input_ids)
answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
print(f"问题: {question}")
print(f"RAG 答案: {answer}")

# 自定义 RAG 系统（LangChain 风格）
from dataclasses import dataclass
from typing import List, Optional

@dataclass
class Document:
    content: str
    metadata: dict

class CustomRAG:
    """自定义 RAG 问答系统"""
    
    def __init__(self, retriever, generator):
        """
        retriever: 检索器 (BM25 / FAISS / embedding-based)
        generator: 生成模型 (LLM API / 本地模型)
        """
        self.retriever = retriever
        self.generator = generator
    
    def answer(self, question: str, top_k: int = 3) -> dict:
        """
        RAG 三步流程：
        1. Retrieve: 检索相关文档
        2. Augment: 构建增强提示词
        3. Generate: 生成答案
        """
        # 步骤 1: 检索
        docs: List[Document] = self.retriever.retrieve(question, k=top_k)
        
        # 步骤 2: 构建增强上下文
        context_parts = []
        for i, doc in enumerate(docs):
            context_parts.append(f"[Document {i+1}] {doc.content}")
        context = "\n\n".join(context_parts)
        
        # 步骤 3: 生成答案
        prompt = self._build_prompt(question, context)
        answer = self.generator.generate(prompt)
        
        return {
            "question": question,
            "answer": answer,
            "source_documents": [
                {"content": d.content[:200], "metadata": d.metadata}
                for d in docs
            ],
        }
    
    def _build_prompt(self, question: str, context: str) -> str:
        """构建 RAG 提示词模板"""
        return f"""基于以下参考文档回答问题。如果文档中没有足够信息，请说明无法回答。

参考文档:
{context}

问题: {question}

答案:"""

# 使用示例
print("=== 自定义 RAG 系统 ===")
print("步骤 1: 构建文档索引 (FAISS / BM25)")
print("步骤 2: 配置生成模型 (LLM)")
print("步骤 3: 调用 rag.answer(question)")
print("\n优势: 答案可追溯、降低幻觉、支持实时更新知识库")

# HotpotQA 多跳推理：管道式方法
from collections import defaultdict

class PipelineMultiHopQA:
    """管道式多跳推理 QA"""
    
    def __init__(self, retriever, reader):
        self.retriever = retriever
        self.reader = reader
    
    def answer_multihop(self, question: str) -> dict:
        """
        多跳推理流程:
        跳 1: 问题 → 检索文档 → 提取中间实体
        跳 2: 中间实体 → 检索新文档 → 提取答案
        """
        # 第一跳：检索与问题直接相关的文档
        hop1_docs = self.retriever.retrieve(question, k=5)
        
        # 从第一跳文档中提取中间实体
        intermediate_entities = []
        for doc in hop1_docs:
            # 使用 NER 或关键词提取中间实体
            entities = self._extract_entities(question, doc)
            intermediate_entities.extend(entities)
        
        # 第二跳：用中间实体检索更多文档
        hop2_docs = []
        for entity in intermediate_entities[:3]:  # 最多取 3 个实体
            related_docs = self.retriever.retrieve(entity, k=3)
            hop2_docs.extend(related_docs)
        
        # 综合所有文档，生成答案
        all_docs = hop1_docs + hop2_docs
        answer, supporting_facts = self.reader.multidoc_read(question, all_docs)
        
        return {
            "answer": answer,
            "supporting_facts": supporting_facts,
            "hop1_docs": [d.title for d in hop1_docs],
            "hop2_docs": [d.title for d in hop2_docs],
            "intermediate_entities": intermediate_entities,
        }
    
    def _extract_entities(self, question, doc):
        """从文档中提取中间实体（简化版）"""
        return ["Entity1", "Entity2"]

# HotpotQA 问题类型
print("=== HotpotQA 问题类型 ===")
types = {
    "bridge": {"描述": "桥接推理", "示例": "A 的 B 是哪个国家的？"},
    "comparison": {"描述": "比较推理", "示例": "A 和 B 谁出生得更早？"},
    "inference": {"描述": "隐含推理", "示例": "A 的导师的学生是谁？"},
}
for t, info in types.items():
    print(f"  {t}: {info['描述']} - {info['示例']}")

# 文档-实体图构建与推理
import torch
import torch.nn as nn
import torch.nn.functional as F

class EntityGraph:
    """文档-实体图"""
    
    def __init__(self):
        self.nodes = []       # 节点列表 (文档节点 + 实体节点)
        self.edges = []       # 边列表 (实体共现关系)
        self.adj_matrix = None  # 邻接矩阵
    
    def build_from_documents(self, documents):
        """
        从文档集合构建图:
        - 每个文档作为一个节点
        - 每个命名实体作为一个节点
        - 如果两个实体在同一文档中出现，则添加边
        """
        node_id = 0
        doc_nodes = []
        entity_nodes = {}
        
        for doc in documents:
            # 添加文档节点
            doc_nodes.append(node_id)
            self.nodes.append({"type": "document", "content": doc.content})
            node_id += 1
            
            # 添加实体节点和边
            for entity in doc.entities:
                if entity not in entity_nodes:
                    entity_nodes[entity] = node_id
                    self.nodes.append({"type": "entity", "name": entity})
                    node_id += 1
                
                # 文档-实体边
                self.edges.append((doc_nodes[-1], entity_nodes[entity]))
        
        # 构建邻接矩阵
        n = len(self.nodes)
        self.adj_matrix = torch.zeros(n, n)
        for i, j in self.edges:
            self.adj_matrix[i][j] = 1
            self.adj_matrix[j][i] = 1  # 无向图
        
        print(f"图构建完成: {n} 个节点, {len(self.edges)} 条边")

class GraphAttentionNetwork(nn.Module):
    """图注意力网络 (GAT) 用于多跳推理"""
    
    def __init__(self, input_dim, hidden_dim, n_heads=4, dropout=0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(
            hidden_dim, n_heads, dropout=dropout, batch_first=True
        )
        self.linear_in = nn.Linear(input_dim, hidden_dim)
        self.linear_out = nn.Linear(hidden_dim, input_dim)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, node_features, adj_matrix):
        """
        node_features: (n_nodes, input_dim) 节点特征
        adj_matrix: (n_nodes, n_nodes) 邻接矩阵
        """
        h = self.linear_in(node_features)
        h = self.dropout(h)
        
        # 使用邻接矩阵作为注意力 mask
        mask = (adj_matrix == 0).bool()
        
        h_attn, _ = self.attention(h.unsqueeze(0), h.unsqueeze(0), h.unsqueeze(0),
                                    key_padding_mask=mask)
        h_attn = h_attn.squeeze(0)
        
        h_out = self.linear_out(h_attn + h)  # 残差连接
        return h_out

print("=== 图推理流程 ===")
print("1. 构建文档-实体图")
print("2. 初始化节点特征 (BERT 编码)")
print("3. GAT 消息传递 (2-3 层)")
print("4. 聚合全局信息 → 回答问题")

# HuggingFace 抽取式 QA 完整训练流程
from datasets import load_dataset
from transformers import (
    AutoTokenizer, AutoModelForQuestionAnswering,
    TrainingArguments, Trainer, DataCollatorWithPadding,
    default_data_collator
)
import numpy as np

# 1. 加载数据和模型
dataset = load_dataset("squad")
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

# 2. 预处理函数
max_length = 384
stride = 128

def preprocess_function(examples):
    questions = [q.strip() for q in examples["question"]]
    inputs = tokenizer(
        questions,
        examples["context"],
        max_length=max_length,
        truncation="only_second",
        stride=stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )
    
    sample_map = inputs.pop("overflow_to_sample_mapping")
    answers = examples["answers"]
    start_positions = []
    end_positions = []
    
    for i, offset in enumerate(inputs["offset_mapping"]):
        sample_idx = sample_map[i]
        answer = answers[sample_idx]
        start_char = answer["answer_start"][0]
        end_char = start_char + len(answer["text"][0])
        
        # 判断答案是否在当前窗口内
        seq_ids = inputs.sequence_ids(i)
        context_start = seq_ids.index(1)
        context_end = len(seq_ids) - 1 - seq_ids[::-1].index(1)
        
        if offset[context_start][0] > end_char or offset[context_end][1] < start_char:
            start_positions.append(0)
            end_positions.append(0)
        else:
            idx = context_start
            while idx <= context_end and offset[idx][0] <= start_char:
                idx += 1
            start_positions.append(idx - 1)
            
            idx = context_end
            while idx >= context_start and offset[idx][1] >= end_char:
                idx -= 1
            end_positions.append(idx + 1)
    
    inputs["start_positions"] = start_positions
    inputs["end_positions"] = end_positions
    return inputs

# 3. 应用预处理
tokenized_datasets = dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=dataset["train"].column_names,
)

print("=== 训练配置 ===")
print(f"训练集窗口数: {len(tokenized_datasets['train']):,}")
print(f"验证集窗口数: {len(tokenized_datasets['validation']):,}")

# QA 系统推理部署
from transformers import pipeline
import torch

# 方法 1: 使用 QuestionAnsweringPipeline（最简单）
qa_pipeline = pipeline(
    "question-answering",
    model="deepset/bert-base-cased-squad2",
    device=0 if torch.cuda.is_available() else -1,
)

result = qa_pipeline(
    question="What is the capital of France?",
    context="Paris is the capital and most populous city of France."
)
print(f"Pipeline 结果: {result}")

# 方法 2: 手动推理（更灵活）
from transformers import AutoTokenizer, AutoModelForQuestionAnswering

tokenizer = AutoTokenizer.from_pretrained("deepset/bert-base-cased-squad2")
model = AutoModelForQuestionAnswering.from_pretrained("deepset/bert-base-cased-squad2")
model.eval()

def answer_question(question: str, context: str, max_len: int = 512):
    """手动 QA 推理，支持长上下文窗口聚合"""
    # 编码
    inputs = tokenizer(
        question, context,
        return_tensors="pt",
        max_length=max_len,
        truncation=True,
        padding="max_length"
    )
    
    # 推理
    with torch.no_grad():
        outputs = model(**inputs)
        start_logits = outputs.start_logits
        end_logits = outputs.end_logits
    
    # 解码
    start_idx = torch.argmax(start_logits).item()
    end_idx = torch.argmax(end_logits).item()
    
    answer_tokens = inputs.input_ids[0][start_idx:end_idx + 1]
    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
    
    return answer

# 方法 3: 长上下文多窗口推理
def answer_long_context(question: str, context: str, max_len: int = 512, stride: int = 128):
    """处理超过 max_len 的长上下文"""
    tokens = tokenizer.encode(context, add_special_tokens=False)
    
    best_score = float('-inf')
    best_answer = ""
    
    for i in range(0, len(tokens), max_len - stride):
        window = tokens[i:i + max_len]
        # 添加特殊 token
        window_with_specials = [tokenizer.cls_token_id] + tokenizer.encode(question) +                                [tokenizer.sep_token_id] + window + [tokenizer.sep_token_id]
        
        if len(window_with_specials) > max_len:
            continue
        
        input_ids = torch.tensor([window_with_specials])
        attention_mask = torch.ones_like(input_ids)
        
        with torch.no_grad():
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        
        # 计算联合概率
        start_probs = torch.softmax(outputs.start_logits, dim=-1)
        end_probs = torch.softmax(outputs.end_logits, dim=-1)
        
        # 在上下文范围内搜索
        question_len = len(tokenizer.encode(question)) + 2
        for s in range(question_len, len(window_with_specials)):
            for e in range(s, min(s + 30, len(window_with_specials))):
                score = start_probs[0][s].item() * end_probs[0][e].item()
                if score > best_score:
                    best_score = score
                    ans_tokens = input_ids[0][s:e+1]
                    best_answer = tokenizer.decode(ans_tokens, skip_special_tokens=True)
    
    return best_answer, best_score

print("=== 部署优化建议 ===")
print("1. 使用 ONNX 导出: python -m transformers.onnx --model=xxx onnx/")
print("2. 使用 ONNX Runtime 加速: pip install onnxruntime-gpu")
print("3. 量化压缩: INT8 量化可减少 75% 模型大小")
print("4. 批处理推理: 合并多个问题的推理请求")