【自然语言处理】基于深度学习基的句子边界检测算法

本文实现了一个基于深度学习基的句子边界检测算法，核心功能是通过 3 种深度学习模型（词嵌入分类器、LSTM+CRF、CNN-LSTM）和 BERT 预训练模型，自动识别文本中句末标点（. ! ?）是否为句子边界，解决普通句末标点、缩写（Mr.）、多段缩写（U.S.A.）、引号内句子等复杂场景的分割问题。代码采用 “数据准备→模型训练→评估→分割应用” 的全流程设计，支持灵活切换模型，适配不同性能需求，并用Python代码完整实现。

二、整体架构与核心目标

核心目标

避免简单按标点分割的缺陷（如将Mr. U.S.A.中的.误判为句子结束），通过深度学习模型学习文本语义和序列特征，精准判断 “边界标点”（句子结束）和 “非边界标点”（缩写、多段缩写中间的.）。

整体流程

数据准备与预处理 → 模型选择与初始化 → 模型训练（含损失监控） → 模型评估（F1/分类报告） → 加载最优模型 → 测试文本分割

三、数据准备与预处理模块

1. `SentenceBoundaryDataset` 类（核心数据集）

功能：统一处理文本数据，生成模型可接收的输入格式（token ID、注意力掩码、标签），支持两种模型类型的输入适配：
- 非 BERT 模型（词嵌入分类器、LSTM+CRF、CNN-LSTM）：
  - 分词：通过_tokenize函数用正则提取单词和独立标点（如Mr.拆分为["Mr", "."]）；
  - 词汇映射：将 token 转为词汇表 ID（未知 token 映射为<UNK>，填充为<PAD>）；
  - 标签构建：生成与 token 长度一致的标签序列（1 = 边界，0 = 非边界），填充位置标注为-100（PyTorch 默认忽略该索引的损失计算）。
- BERT 类模型（bert/bert_crf）：
  - 用 HuggingFace 的BertTokenizer处理文本，添加[CLS] [SEP]特殊 token，padding/truncation 到固定长度（128）；
  - 标签调整：仅保留标点 token 的标签（非标点 token 标注为-100），确保损失计算仅针对标点位置。
关键细节：标签中的-100用于过滤非标点 token 和填充 token，避免无效数据影响模型训练。

2. `prepare_data` 函数（数据与词汇表准备）

功能：提供带标注的训练数据，构建非 BERT 模型所需的词汇表。
- 标注数据：包含 7 类核心场景（普通句末标点、单段缩写、多段缩写、引号内句子、复杂并列句等），每条数据格式为(文本, [(标点位置, 标签)])，标签1= 边界，0= 非边界；
- 词汇表构建：遍历所有文本的 token，生成{"<PAD>":0, "<UNK>":1, "Mr":2, "." :3,...}格式的词汇表，为非 BERT 模型提供词嵌入映射依据。
输出：拆分后的训练集 / 验证集（8:2 比例）、词汇表（非 BERT 模型用）。

四、模型定义模块（3 种基础模型 + 1 种预训练模型）

代码提供 4 种可切换的模型，覆盖从基础到复杂的不同场景需求，核心差异在于特征提取方式和序列建模能力：

1. `EmbeddingClassifier`（词嵌入 + 滑动窗口分类器）

核心定位：基础轻量模型，适合简单场景（无复杂缩写、短文本）。
结构与功能：
1. 词嵌入层：将 token ID 转为固定维度的词向量（默认 100 维）；
2. 滑动窗口特征提取：用 3 窗口（可配置）捕捉标点前后的上下文特征（如[Mr, ., Smith]），补全边缘窗口的填充；
3. 全连接分类器：通过两层全连接网络（含 ReLU 激活和 Dropout）输出边界 / 非边界的预测概率。
优势：训练快、推理快，无需复杂序列建模；劣势：无法捕捉长距离序列依赖（如多段缩写的连续.）。

2. `LSTMCRF`（LSTM + 条件随机场）

核心定位：复杂场景最优模型，重点解决序列依赖和标签约束（如 “连续边界不可能”）。
结构与功能：
1. 词嵌入层：与上述一致，生成词向量；
2. 双向 LSTM 层：捕捉文本的前后向序列特征（默认 64 维隐藏层，双向拼接为 128 维），处理变长序列（通过pack_padded_sequence）；
3. 线性层：将 LSTM 输出映射为标签发射分数（2 类：边界 / 非边界）；
4. CRF 层（核心）：
  - 学习标签转移概率（如 “边界→非边界” 概率高，“边界→边界” 概率低）；
  - 初始化时强制禁止连续边界（transitions.data[1,1] = -1000.0），符合语言逻辑；
  - 训练时计算对数似然损失，预测时用维特比算法解码最优标签序列。
优势：同时捕捉序列特征和标签依赖，对多段缩写、引号内句子等复杂场景适配性最好；劣势：训练和推理速度略慢于基础模型。

3. `CNNLSTM`（CNN+LSTM 融合模型）

核心定位：平衡局部特征与序列特征，适合中等复杂度文本。
结构与功能：
1. 词嵌入层：生成词向量；
2. CNN 层（1 维卷积）：用 3 核捕捉局部上下文特征（如Fig. 3的局部关联），增强缩写识别能力；
3. 双向 LSTM 层：捕捉全局序列特征；
4. 全连接层：输出预测概率。
优势：CNN 补全 LSTM 的局部特征捕捉短板，LSTM 弥补 CNN 的长序列建模不足，性能均衡；劣势：对极复杂的多段缩写（如U.S.A.）效果略逊于 LSTM+CRF。

4. BERT 预训练模型

核心定位：高精度模型，适合大规模文本或高准确率需求场景。
功能：复用 HuggingFace 的BertForTokenClassification，将句子边界检测视为 token 级分类任务（2 类标签）；
优势：预训练模型自带丰富的语义理解能力，无需手动设计特征，对罕见缩写、复杂句法结构适配性强；劣势：参数量大，训练 / 推理速度慢，需更多计算资源。

五、训练与评估模块

1. `train_model` 函数（模型训练核心）

功能：统一的训练框架，支持 4 种模型的训练，自动处理不同模型的输入格式和损失计算。
核心逻辑：
1. 设备适配：自动检测 GPU/CPU，将模型移至对应设备；
2. 优化器与损失函数：用 Adam 优化器（学习率 0.001），非 BERT 模型用交叉熵损失（忽略-100标签），LSTMCRF 直接输出损失，BERT 模型用自带损失函数；
3. 训练监控：记录训练 / 验证损失，绘制损失曲线，直观观察模型收敛情况；
4. 最佳模型保存：仅保存验证集 F1 最高的模型，避免过拟合。

2. `evaluate_model` 函数（模型评估）

功能：用 F1 分数（加权平均）和分类报告（精确率、召回率、F1）评估模型性能，确保评估结果准确。
关键细节：
- 过滤无效标签：仅统计非-100的有效标签（标点位置），避免填充和非标点 token 影响评估；
- 适配变长序列：对非 BERT 模型，按样本实际长度截取预测结果和标签；
- 输出分类报告：分别展示边界（1）和非边界（0）的性能，明确模型在不同场景的短板（如是否漏判边界、误判缩写）。

六、句子分割模块（`split_sentences` 函数）

功能：将训练好的模型应用于实际文本，根据模型预测的边界位置，分割出独立句子。
核心逻辑：
- BERT 类模型：
  1. 用BertTokenizer编码文本，获取 token 映射；
  2. 模型预测 token 级标签，筛选出 “预测为 1 且是句末标点” 的 token；
  3. 通过token_to_chars将 token 位置映射到原始文本的字符位置，得到边界索引。
- 非 BERT 模型：
  1. 分词后将 token 转为词汇表 ID，输入模型得到预测标签；
  2. 计算每个 token 在原始文本中的字符位置，筛选出 “预测为 1 且是句末标点” 的位置作为边界；
- 句子分割：按边界索引截取文本，过滤空字符串，输出最终分割后的句子列表。

七、主函数 `main`（流程串联）

功能：串联整个系统的流程，从数据准备到最终分割，一键运行。
核心步骤：
1. 数据准备：调用prepare_data获取训练 / 验证集和词汇表；
2. 模型选择：默认选择 LSTMCRF（性能最优），可通过修改model_types切换模型；
3. 数据加载：根据模型类型，生成对应的数据加载器（BERT 模型用BertTokenizer，非 BERT 模型用自定义collate_fn处理变长序列）；
4. 模型训练：调用train_model训练模型，加载最佳模型；
5. 测试与输出：用复杂测试文本（含多种场景）测试模型分割效果，输出分割后的句子。

八、基于深度学习基的句子边界检测算法的Python代码完整实现

import re
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score, classification_report
from transformers import BertTokenizer, BertForTokenClassification
import matplotlib.pyplot as plt

# 设置随机种子，保证结果可复现
torch.manual_seed(42)
np.random.seed(42)


# -------------------------- 1. 数据准备与预处理 --------------------------
class SentenceBoundaryDataset(Dataset):
    """句子边界检测数据集"""

    def __init__(self, texts, labels=None, tokenizer=None, max_len=128, model_type="lstm"):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.model_type = model_type

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]

        if self.model_type in ["bert", "bert_crf"]:
            # BERT模型处理
            encoding = self.tokenizer(
                text,
                add_special_tokens=True,
                max_length=self.max_len,
                return_token_type_ids=False,
                padding="max_length",
                truncation=True,
                return_attention_mask=True,
                return_tensors="pt"
            )

            item = {
                "input_ids": encoding["input_ids"].flatten(),
                "attention_mask": encoding["attention_mask"].flatten()
            }

            if self.labels is not None:
                # 调整标签长度以匹配tokenizer输出
                labels = self.labels[idx]
                adjusted_labels = [-100]  # CLS token
                ptr = 0

                for token in self.tokenizer.tokenize(text):
                    if ptr < len(labels) and token == labels[ptr][0]:
                        adjusted_labels.append(labels[ptr][1])
                        ptr += 1
                    else:
                        adjusted_labels.append(-100)  # 非标点token不参与损失计算

                adjusted_labels.append(-100)  # SEP token
                adjusted_labels = adjusted_labels[:self.max_len]
                if len(adjusted_labels) < self.max_len:
                    adjusted_labels += [-100] * (self.max_len - len(adjusted_labels))

                item["labels"] = torch.tensor(adjusted_labels, dtype=torch.long)

            return item

        else:
            # 其他模型处理
            tokens = self._tokenize(text)
            token_ids = [self.tokenizer.get(token, self.tokenizer["<UNK>"]) for token in tokens]

            item = {
                "tokens": torch.tensor(token_ids, dtype=torch.long),
                "length": torch.tensor(len(token_ids), dtype=torch.long)
            }

            if self.labels is not None:
                # 创建标签序列，1表示边界，0表示非边界
                labels = torch.zeros(len(tokens), dtype=torch.long)
                for pos, label in self.labels[idx]:
                    if pos < len(tokens):
                        labels[pos] = label
                item["labels"] = labels

            return item

    def _tokenize(self, text):
        """简单分词，保留标点符号作为独立token"""
        return re.findall(r"\w+|[^\w\s]", text)


def prepare_data():
    """准备训练数据"""
    # 带标注的训练数据：(文本, [(标点位置, 标签)])
    # 标签说明：1=句子边界，0=非边界
    labeled_data = [
        # 普通句末标点
        ("He went to school. She stayed home.", [(8, 1), (17, 1)]),
        ("I love reading. It broadens my horizon!", [(12, 1), (31, 1)]),
        ("Where are you going?", [(18, 1)]),

        # 单段缩写（非边界）
        ("Mr. Smith came to the party.", [(2, 0), (23, 1)]),
        ("Mrs. Brown is our new teacher.", [(3, 0), (28, 1)]),
        ("Dr. Wang published a paper.", [(2, 0), (23, 1)]),
        ("Fig. 2 shows the result.", [(3, 0), (21, 1)]),
        ("Eq. 5 is derived.", [(2, 0), (16, 1)]),

        # 多段缩写（非边界）
        ("U.S.A. is powerful.", [(1, 0), (3, 0), (5, 0), (14, 1)]),
        ("e.g. apple is fruit.", [(1, 0), (19, 1)]),
        ("Ph.D. student won.", [(1, 0), (3, 0), (16, 1)]),
        ("N.Y.C. is big.", [(1, 0), (3, 0), (10, 1)]),

        # 引号内场景
        ("He said, \"I'm done.\" She smiled.", [(13, 1), (21, 1), (32, 1)]),
        ("She shouted, \"Help!\"", [(14, 1), (20, 1)]),
        ("\"Hello!\" He waved.", [(6, 1), (14, 1)]),

        # 复杂场景
        ("She bought milk, bread, etc. and went home.", [(25, 0), (42, 1)]),
        ("Eq. 2 and Fig. 3 are referenced.", [(2, 0), (13, 0), (31, 1)]),
    ]

    # 拆分文本和标签
    texts = [item[0] for item in labeled_data]
    labels = [item[1] for item in labeled_data]

    # 创建词汇表（用于非BERT模型）
    tokenizer = {"<PAD>": 0, "<UNK>": 1}
    for text in texts:
        tokens = re.findall(r"\w+|[^\w\s]", text)
        for token in tokens:
            if token not in tokenizer:
                tokenizer[token] = len(tokenizer)

    return train_test_split(texts, labels, test_size=0.2, random_state=42), tokenizer


# -------------------------- 2. 模型定义 --------------------------

class EmbeddingClassifier(nn.Module):
    """词嵌入 + 分类器模型"""

    def __init__(self, vocab_size, embedding_dim=100, hidden_dim=64, window_size=3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.window_size = window_size
        self.fc1 = nn.Linear(embedding_dim * window_size, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, 2)  # 二分类：边界/非边界
        self.activation = nn.ReLU()
        self.dropout = nn.Dropout(0.3)

    def forward(self, x, lengths=None):
        # x shape: (batch_size, seq_len)
        batch_size, seq_len = x.shape

        # 嵌入层
        embedded = self.embedding(x)  # (batch_size, seq_len, embedding_dim)

        # 滑动窗口特征提取
        windows = []
        for i in range(seq_len):
            start = max(0, i - self.window_size // 2)
            end = min(seq_len, i + self.window_size // 2 + 1)

            # 补全窗口大小
            if end - start < self.window_size:
                if start == 0:
                    pad = torch.zeros(batch_size, self.window_size - (end - start), embedded.shape[2],
                                      device=embedded.device)
                    window = torch.cat([embedded[:, start:end, :], pad], dim=1)
                else:
                    pad = torch.zeros(batch_size, self.window_size - (end - start), embedded.shape[2],
                                      device=embedded.device)
                    window = torch.cat([pad, embedded[:, start:end, :]], dim=1)
            else:
                window = embedded[:, start:end, :]

            windows.append(window)

        # 拼接所有窗口
        window_features = torch.stack(windows, dim=1)  # (batch_size, seq_len, window_size, embedding_dim)
        window_features = window_features.view(batch_size, seq_len, -1)  # 展平

        # 分类器
        out = self.fc1(window_features)
        out = self.activation(out)
        out = self.dropout(out)
        out = self.fc2(out)  # (batch_size, seq_len, 2)

        return out


class LSTMCRF(nn.Module):
    """LSTM + CRF模型"""

    def __init__(self, vocab_size, embedding_dim=100, hidden_dim=64, num_tags=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            input_size=embedding_dim,
            hidden_size=hidden_dim // 2,  # 双向
            bidirectional=True,
            batch_first=True
        )
        self.fc = nn.Linear(hidden_dim, num_tags)
        self.num_tags = num_tags
        self.pad_idx = -100  # 填充标签索引

        # CRF层参数
        self.transitions = nn.Parameter(torch.randn(num_tags, num_tags))
        self.start_transitions = nn.Parameter(torch.randn(num_tags))
        self.end_transitions = nn.Parameter(torch.randn(num_tags))

        # 初始化转移矩阵：禁止不可能的转移（如边界→边界连续）
        self.transitions.data[1, 1] = -1000.0  # 1=边界标签，禁止连续边界

    def forward(self, x, lengths=None, tags=None):
        # x shape: (batch_size, seq_len)
        embedded = self.embedding(x)  # (batch_size, seq_len, embedding_dim)

        # 处理变长序列
        if lengths is not None:
            packed = pack_padded_sequence(embedded, lengths, batch_first=True, enforce_sorted=False)
            lstm_out, _ = self.lstm(packed)
            lstm_out, _ = pad_packed_sequence(lstm_out, batch_first=True)
        else:
            lstm_out, _ = self.lstm(embedded)

        # 线性层得到发射分数
        emissions = self.fc(lstm_out)  # (batch_size, seq_len, num_tags)

        if tags is None:
            # 预测模式：维特比解码
            return self._viterbi_decode(emissions, lengths)
        else:
            # 训练模式：计算损失（过滤填充标签）
            loss = self._crf_loss(emissions, tags, lengths)
            return loss

    def _crf_loss(self, emissions, tags, lengths):
        """计算CRF损失"""
        batch_size, seq_len = emissions.shape[:2]

        # 过滤填充标签（只保留有效标签位置）
        mask = (tags != self.pad_idx).float()  # (batch_size, seq_len)
        valid_tags = tags.masked_fill(tags == self.pad_idx, 0)  # 填充位置临时设为0（不影响计算）

        # 1. 计算有效路径分数
        # 初始分数：start_transitions + 第一个有效时间步的发射分数
        start_mask = (lengths > 0).float()  # (batch_size,)
        first_tags = valid_tags[:, 0]  # (batch_size,)
        total_score = self.start_transitions[first_tags] * start_mask

        # 第一个有效时间步的发射分数
        first_emissions = emissions[:, 0].gather(1, first_tags.unsqueeze(1)).squeeze(1)  # (batch_size,)
        total_score += first_emissions * start_mask

        # 迭代计算后续时间步分数
        for i in range(1, seq_len):
            # 当前时间步有效掩码（样本未结束且当前位置非填充）
            step_mask = (i < lengths).float() * mask[:, i]  # (batch_size,)
            if step_mask.sum() == 0:
                continue  # 无有效样本，跳过

            # 前一个时间步标签和当前时间步标签
            prev_tags = valid_tags[:, i - 1]  # (batch_size,)
            curr_tags = valid_tags[:, i]  # (batch_size,)

            # 转移分数：prev_tags → curr_tags
            trans_score = self.transitions[prev_tags, curr_tags]  # (batch_size,)

            # 当前时间步发射分数
            emit_score = emissions[:, i].gather(1, curr_tags.unsqueeze(1)).squeeze(1)  # (batch_size,)

            # 累加分数（只累加有效样本）
            total_score += (trans_score + emit_score) * step_mask

        # 加上结束转移分数
        last_valid_idx = (lengths - 1).clamp(min=0)  # (batch_size,)
        last_tags = valid_tags.gather(1, last_valid_idx.unsqueeze(1)).squeeze(1)  # (batch_size,)
        end_score = self.end_transitions[last_tags] * start_mask
        total_score += end_score

        # 2. 计算配分函数（归一化因子）
        log_partition = self._forward_algorithm(emissions, lengths, mask)

        # 3. 损失 = 负的平均对数似然
        return (log_partition - total_score).sum() / batch_size

    def _forward_algorithm(self, emissions, lengths, mask):
        """前向算法计算配分函数（适配变长序列）"""
        batch_size, seq_len, num_tags = emissions.shape

        # 初始化分数：start_transitions + 第一个时间步发射分数
        log_probs = self.start_transitions.unsqueeze(0) + emissions[:, 0]  # (batch_size, num_tags)

        # 应用第一个时间步掩码
        first_mask = mask[:, 0].unsqueeze(1)  # (batch_size, 1)
        log_probs = log_probs * first_mask + (-1e18) * (1 - first_mask)

        # 迭代计算后续时间步
        for i in range(1, seq_len):
            # 当前时间步掩码
            step_mask = mask[:, i].unsqueeze(1)  # (batch_size, 1)
            if step_mask.sum() == 0:
                continue

            # 广播机制：(batch, num_tags, 1) + (num_tags, num_tags) + (batch, 1, num_tags)
            next_log_probs = log_probs.unsqueeze(2) + self.transitions.unsqueeze(0) + emissions[:, i].unsqueeze(1)

            # 对数求和exp
            next_log_probs = torch.logsumexp(next_log_probs, dim=1)  # (batch_size, num_tags)

            # 应用掩码：有效位置更新，无效位置保持原分数
            log_probs = next_log_probs * step_mask + log_probs * (1 - step_mask)

        # 加上结束转移分数并求和
        log_probs += self.end_transitions.unsqueeze(0)
        return torch.logsumexp(log_probs, dim=1).sum()  # 累加所有样本的配分函数

    def _viterbi_decode(self, emissions, lengths):
        """维特比算法解码最优路径（适配变长序列）"""
        batch_size, seq_len, num_tags = emissions.shape

        # 初始化分数和路径
        scores = self.start_transitions.unsqueeze(0) + emissions[:, 0]  # (batch_size, num_tags)
        paths = []

        # 迭代计算每个时间步
        for i in range(1, seq_len):
            # 计算所有可能转移的分数
            next_scores = scores.unsqueeze(2) + self.transitions.unsqueeze(0) + emissions[:, i].unsqueeze(1)

            # 记录最佳前序标签和最大分数
            max_scores, best_prev_tags = next_scores.max(dim=1)  # (batch_size, num_tags)
            paths.append(best_prev_tags)
            scores = max_scores

        # 加上结束转移分数并找到最佳最终标签
        scores += self.end_transitions.unsqueeze(0)
        best_tags = scores.argmax(dim=1)  # (batch_size,)

        # 回溯找到最佳路径
        best_paths = [best_tags.unsqueeze(1)]
        for i in reversed(range(seq_len - 1)):
            best_tags = paths[i].gather(1, best_tags.unsqueeze(1)).squeeze(1)
            best_paths.insert(0, best_tags.unsqueeze(1))

        best_paths = torch.cat(best_paths, dim=1)  # (batch_size, seq_len)

        # 适配变长序列：超过长度的位置设为填充标签
        for i in range(batch_size):
            best_paths[i, lengths[i]:] = self.pad_idx

        return best_paths


class CNNLSTM(nn.Module):
    """CNN-LSTM融合模型"""

    def __init__(self, vocab_size, embedding_dim=100, cnn_out_dim=64, lstm_hidden_dim=64, num_tags=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)

        # CNN层：捕捉局部特征
        self.cnn = nn.Conv1d(
            in_channels=embedding_dim,
            out_channels=cnn_out_dim,
            kernel_size=3,
            padding=1
        )

        # LSTM层：捕捉序列特征
        self.lstm = nn.LSTM(
            input_size=cnn_out_dim,
            hidden_size=lstm_hidden_dim // 2,
            bidirectional=True,
            batch_first=True
        )

        # 输出层
        self.fc = nn.Linear(lstm_hidden_dim, num_tags)
        self.dropout = nn.Dropout(0.3)

    def forward(self, x, lengths=None):
        # x shape: (batch_size, seq_len)
        embedded = self.embedding(x)  # (batch_size, seq_len, embedding_dim)

        # CNN层需要转置维度
        cnn_out = self.cnn(embedded.transpose(1, 2)).transpose(1, 2)  # (batch_size, seq_len, cnn_out_dim)
        cnn_out = self.dropout(cnn_out)

        # LSTM层
        if lengths is not None:
            packed = pack_padded_sequence(cnn_out, lengths, batch_first=True, enforce_sorted=False)
            lstm_out, _ = self.lstm(packed)
            lstm_out, _ = pad_packed_sequence(lstm_out, batch_first=True)
        else:
            lstm_out, _ = self.lstm(cnn_out)

        # 输出层
        out = self.fc(lstm_out)  # (batch_size, seq_len, num_tags)
        return out


# -------------------------- 3. 训练与评估函数 --------------------------

def train_model(model, train_loader, val_loader, model_type, epochs=20, lr=0.001):
    """训练模型"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss(ignore_index=-100)

    train_losses = []
    val_losses = []
    best_f1 = 0.0

    for epoch in range(epochs):
        model.train()
        train_loss = 0.0

        for batch in train_loader:
            optimizer.zero_grad()

            if model_type == "bert" or model_type == "bert_crf":
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                labels = batch["labels"].to(device)

                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                loss = outputs.loss
            else:
                tokens = batch["tokens"].to(device)
                lengths = batch["length"].to(device)
                labels = batch["labels"].to(device)

                if model_type == "lstm_crf":
                    loss = model(tokens, lengths, labels)
                else:
                    outputs = model(tokens, lengths)
                    loss = criterion(outputs.transpose(1, 2), labels)

            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        train_loss /= len(train_loader)
        train_losses.append(train_loss)

        # 验证
        val_loss, val_f1 = evaluate_model(model, val_loader, model_type, criterion)
        val_losses.append(val_loss)

        print(f"Epoch {epoch + 1}/{epochs}")
        print(f"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val F1: {val_f1:.4f}")

        # 保存最佳模型
        if val_f1 > best_f1:
            best_f1 = val_f1
            torch.save(model.state_dict(), f"best_{model_type}_model.pt")

    # 绘制损失曲线
    plt.plot(train_losses, label="Train Loss")
    plt.plot(val_losses, label="Val Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    plt.title(f"Training and Validation Loss ({model_type})")
    plt.show()

    return model


def evaluate_model(model, dataloader, model_type, criterion=None):
    """评估模型"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    total_loss = 0.0
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for batch in dataloader:
            if model_type == "bert" or model_type == "bert_crf":
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                labels = batch["labels"].to(device)

                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                logits = outputs.logits if model_type == "bert" else outputs
                loss = outputs.loss

                # 获取预测结果
                preds = torch.argmax(logits, dim=2)

                # 收集非忽略的标签和预测
                mask = (labels != -100)
                all_preds.extend(preds[mask].cpu().numpy())
                all_labels.extend(labels[mask].cpu().numpy())

                total_loss += loss.item()
            else:
                tokens = batch["tokens"].to(device)
                lengths = batch["length"].to(device)
                labels = batch["labels"].to(device)

                if model_type == "lstm_crf":
                    preds = model(tokens, lengths)
                    loss = model(tokens, lengths, labels)
                else:
                    outputs = model(tokens, lengths)
                    preds = torch.argmax(outputs, dim=2)
                    loss = criterion(outputs.transpose(1, 2), labels)

                # 收集结果（排除填充部分）
                for i in range(len(lengths)):
                    length = lengths[i].item()
                    # 过滤填充标签
                    valid_mask = (labels[i, :length] != -100)
                    all_preds.extend(preds[i, :length][valid_mask].cpu().numpy())
                    all_labels.extend(labels[i, :length][valid_mask].cpu().numpy())

                total_loss += loss.item()

    avg_loss = total_loss / len(dataloader)
    f1 = f1_score(all_labels, all_preds, average="weighted")

    # 详细分类报告
    print(classification_report(all_labels, all_preds, zero_division=0))

    return avg_loss, f1


def split_sentences(text, model, tokenizer, model_type, device):
    """使用训练好的模型分割句子"""
    model.eval()
    model.to(device)

    if model_type in ["bert", "bert_crf"]:
        # BERT模型处理
        encoding = tokenizer(
            text,
            add_special_tokens=True,
            return_token_type_ids=False,
            return_attention_mask=True,
            return_tensors="pt"
        )

        input_ids = encoding["input_ids"].to(device)
        attention_mask = encoding["attention_mask"].to(device)

        with torch.no_grad():
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            logits = outputs.logits
            preds = torch.argmax(logits, dim=2).squeeze().cpu().numpy()

        # 将token映射回原始文本
        tokens = tokenizer.convert_ids_to_tokens(input_ids.squeeze().cpu().numpy())
        boundaries = []

        for i, (token, pred) in enumerate(zip(tokens, preds)):
            if pred == 1 and token in ['.', '!', '?']:
                # 找到在原始文本中的位置
                char_span = tokenizer.token_to_chars(0, i)
                if char_span:
                    boundaries.append(char_span.end)

    else:
        # 其他模型处理
        tokens = re.findall(r"\w+|[^\w\s]", text)
        token_ids = [tokenizer.get(token, tokenizer["<UNK>"]) for token in tokens]
        token_tensor = torch.tensor([token_ids], dtype=torch.long).to(device)
        length_tensor = torch.tensor([len(token_ids)], dtype=torch.long).to(device)

        with torch.no_grad():
            if model_type == "lstm_crf":
                preds = model(token_tensor, length_tensor).squeeze().cpu().numpy()
            else:
                outputs = model(token_tensor, length_tensor)
                preds = torch.argmax(outputs, dim=2).squeeze().cpu().numpy()

        # 找到边界位置
        boundaries = []
        current_pos = 0
        for i, (token, pred) in enumerate(zip(tokens, preds)):
            # 跳过填充标签
            if pred == -100:
                current_pos += len(token) + (1 if i < len(tokens) - 1 else 0)
                continue
            current_pos += len(token) + (1 if i < len(tokens) - 1 else 0)  # +1 是空格
            if pred == 1 and token in ['.', '!', '?']:
                boundaries.append(current_pos)

    # 根据边界分割句子
    sentences = []
    start = 0
    for boundary in sorted(boundaries):
        sentences.append(text[start:boundary].strip())
        start = boundary

    # 添加最后一句
    if start < len(text):
        sentences.append(text[start:].strip())

    return [s for s in sentences if s]  # 过滤空字符串


# -------------------------- 4. 主函数 --------------------------

def main():
    # 准备数据
    (train_texts, val_texts, train_labels, val_labels), tokenizer = prepare_data()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")

    # 模型类型选择
    model_types = ["embedding_classifier", "lstm_crf", "cnn_lstm", "bert"]
    model_type = model_types[1]  # 默认使用LSTM+CRF模型
    print(f"训练模型: {model_type}")

    # 初始化数据加载器
    if model_type in ["bert", "bert_crf"]:
        # BERT模型
        bert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

        train_dataset = SentenceBoundaryDataset(
            train_texts, train_labels, bert_tokenizer, model_type=model_type
        )
        val_dataset = SentenceBoundaryDataset(
            val_texts, val_labels, bert_tokenizer, model_type=model_type
        )

        train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=4)

        # 初始化BERT模型
        if model_type == "bert":
            model = BertForTokenClassification.from_pretrained(
                "bert-base-uncased",
                num_labels=2,
                ignore_mismatched_sizes=True
            )
        else:  # bert_crf（简化版）
            model = BertForTokenClassification.from_pretrained(
                "bert-base-uncased",
                num_labels=2,
                ignore_mismatched_sizes=True
            )

    else:
        # 其他模型
        vocab_size = len(tokenizer)

        train_dataset = SentenceBoundaryDataset(
            train_texts, train_labels, tokenizer, model_type=model_type
        )
        val_dataset = SentenceBoundaryDataset(
            val_texts, val_labels, tokenizer, model_type=model_type
        )

        # 自定义collate_fn处理变长序列
        def collate_fn(batch):
            tokens = [item["tokens"] for item in batch]
            lengths = [item["length"] for item in batch]
            labels = [item["labels"] for item in batch]

            # 填充序列
            tokens_padded = pad_sequence(tokens, batch_first=True, padding_value=0)
            # 标签填充用-100（忽略损失计算）
            labels_padded = pad_sequence(labels, batch_first=True, padding_value=-100)

            return {
                "tokens": tokens_padded,
                "length": torch.tensor(lengths),
                "labels": labels_padded
            }

        train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=collate_fn)
        val_loader = DataLoader(val_dataset, batch_size=4, collate_fn=collate_fn)

        # 初始化模型
        if model_type == "embedding_classifier":
            model = EmbeddingClassifier(vocab_size)
        elif model_type == "lstm_crf":
            model = LSTMCRF(vocab_size)
        elif model_type == "cnn_lstm":
            model = CNNLSTM(vocab_size)

    # 训练模型
    model = train_model(model, train_loader, val_loader, model_type)

    # 加载最佳模型
    model.load_state_dict(torch.load(f"best_{model_type}_model.pt", map_location=device))

    # 测试分割效果
    test_text = """Mr. Smith went to Dr. Lee's office. They discussed Fig. 3 and Eq. 2. 
    U.S.A. has a long history. etc. is often used in academic papers. He said, "I'm busy!" She nodded.
    Dr. Wang published a paper in 2024. It references Eq. 5 and Fig. 7. 
    e.g. apple, banana and orange are fruits. N.Y.C. is a big city. etc. should be used carefully. Where are you going?"""

    print("\n=== 测试文本 ===")
    print(test_text)

    # 分割句子
    if model_type in ["bert", "bert_crf"]:
        sentences = split_sentences(test_text, model, bert_tokenizer, model_type, device)
    else:
        sentences = split_sentences(test_text, model, tokenizer, model_type, device)

    print("\n=== 分割结果 ===")
    for i, sent in enumerate(sentences, 1):
        print(f"{i}. {sent}")


if __name__ == "__main__":
    main()