基于 BERT 的抽取式摘要

pip install torch  # 或 tensorflow，取决于你选择哪个深度学习框架
pip install transformers  # Hugging Face 的 Transformers 库
pip install jieba  # 中文分词
pip install scikit-learn  # 机器学习工具库
pip install rouge-score # ROUGE 评估 (或使用 py-rouge)
# 如果使用 py-rouge (更完整的 ROUGE 实现):
# pip install py-rouge  # 可能需要先安装 Perl

import json
import jieba
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer

class SummarizationDataset(Dataset):
    def __init__(self, data_file, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = self.load_data(data_file)

    def load_data(self, data_file):
        data = []
        with open(data_file, 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line.strip())
                # text = item['text']
                text = item['title']  #根据具体情况选择title 还是 content
                # summary = item['summary']
                summary = item['content']
                data.append((text, summary))
        return data

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

    def __getitem__(self, idx):
        text, summary = self.data[idx]
        # 分词并添加特殊标记
        text_tokens = ['[CLS]'] + list(jieba.cut(text)) + ['[SEP]']
        summary_tokens = ['[CLS]'] + list(jieba.cut(summary)) + ['[SEP]']

        # 构建标签（抽取式摘要：0 或 1）
        labels = [0] * len(text_tokens)
        summary_token_ids = set(self.tokenizer.convert_tokens_to_ids(summary_tokens))
        for i, token in enumerate(text_tokens):
            if self.tokenizer.convert_tokens_to_ids(token) in summary_token_ids:
                labels[i] = 1 #如果需要更精细的label, 可以计算每个句子和summary的rouge值作为label

        # 截断或填充到最大长度
        text_tokens = text_tokens[:self.max_length]
        labels = labels[:self.max_length]

        text_ids = self.tokenizer.convert_tokens_to_ids(text_tokens)
        attention_mask = [1] * len(text_ids)

        padding_length = self.max_length - len(text_ids)
        text_ids += [self.tokenizer.pad_token_id] * padding_length
        attention_mask += [0] * padding_length
        labels += [0] * padding_length  # 标签也需要填充

        return {
            'input_ids': torch.tensor(text_ids, dtype=torch.long),
            'attention_mask': torch.tensor(attention_mask, dtype=torch.long),
            'labels': torch.tensor(labels, dtype=torch.long)
        }

# 使用示例
if __name__ == '__main__':
    tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
    dataset = SummarizationDataset('train.txt', tokenizer)  # 替换为你的训练数据文件
    dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

    for batch in dataloader:
        print(batch['input_ids'].shape)
        print(batch['attention_mask'].shape)
        print(batch['labels'].shape)
        break

import torch
import torch.nn as nn
from transformers import BertModel

class BertForExtractiveSummarization(nn.Module):
    def __init__(self, bert_model_name='bert-base-chinese'):
        super(BertForExtractiveSummarization, self).__init__()
        self.bert = BertModel.from_pretrained(bert_model_name)
        self.classifier = nn.Linear(self.bert.config.hidden_size, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        # 使用 [CLS] 标记的输出，或者所有 token 输出的平均池化
        # pooled_output = outputs.pooler_output  # [CLS] 标记
        last_hidden_state = outputs.last_hidden_state # 取最后一个hidden_state
        logits = self.classifier(last_hidden_state) #
        probs = self.sigmoid(logits).squeeze(-1)  # 转换为概率
        return probs

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from transformers import BertTokenizer
from model import BertForExtractiveSummarization  # 从 model.py 导入模型
from data_utils import SummarizationDataset  # 从 data_utils.py 导入数据集
from tqdm import tqdm  # 导入 tqdm

# 超参数
BATCH_SIZE = 8
LEARNING_RATE = 2e-5
NUM_EPOCHS = 3
MAX_LENGTH = 512
BERT_MODEL_NAME = 'bert-base-chinese'

# 设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载 tokenizer 和数据集
tokenizer = BertTokenizer.from_pretrained(BERT_MODEL_NAME)
train_dataset = SummarizationDataset('train.txt', tokenizer, MAX_LENGTH)  # 替换为你的训练数据文件
train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

dev_dataset = SummarizationDataset('dev.txt', tokenizer, MAX_LENGTH)  # 替换为你的验证数据文件
dev_dataloader = DataLoader(dev_dataset, batch_size=BATCH_SIZE, shuffle=False)


# 初始化模型、优化器和损失函数
model = BertForExtractiveSummarization(BERT_MODEL_NAME).to(device)
optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE)  # 使用 AdamW 优化器
criterion = nn.BCELoss()  # 二元交叉熵损失

# 训练循环
for epoch in range(NUM_EPOCHS):
    model.train()  # 设置为训练模式
    total_loss = 0
    for batch in tqdm(train_dataloader, desc=f"Epoch {epoch + 1}/{NUM_EPOCHS}"):
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()  # 清空梯度

        probs = model(input_ids, attention_mask)
        loss = criterion(probs, labels.float())

        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数

        total_loss += loss.item()

    print(f"Epoch {epoch + 1}/{NUM_EPOCHS}, Loss: {total_loss / len(train_dataloader)}")

    # 验证 (可选，但强烈建议)
    model.eval()  # 设置为评估模式
    with torch.no_grad():  # 关闭梯度计算
        val_loss = 0
        for batch in dev_dataloader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            probs = model(input_ids, attention_mask)
            loss = criterion(probs, labels.float())

            val_loss += loss.item()

        print(f"Validation Loss: {val_loss / len(dev_dataloader)}")

# 保存模型
torch.save(model.state_dict(), 'extractive_summarizer.pth')

import torch
from transformers import BertTokenizer
from model import BertForExtractiveSummarization  # 导入模型
import jieba

def summarize(text, model, tokenizer, max_length=512, threshold=0.5):
    """
    使用训练好的模型进行抽取式摘要。

    Args:
        text: 要摘要的文本。
        model: 训练好的模型。
        tokenizer: 分词器。
        max_length: 最大序列长度。
        threshold: 概率阈值，用于决定是否选择句子。

    Returns:
        str: 抽取式摘要。
    """
    model.eval()  # 设置为评估模式
    tokens = ['[CLS]'] + list(jieba.cut(text)) + ['[SEP]']
    tokens = tokens[:max_length]
    input_ids = tokenizer.convert_tokens_to_ids(tokens)
    attention_mask = [1] * len(input_ids)
    padding_length = max_length - len(input_ids)
    input_ids += [tokenizer.pad_token_id] * padding_length
    attention_mask += [0] * padding_length

    input_ids = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0).to(device)  # 添加 batch 维度
    attention_mask = torch.tensor(attention_mask, dtype=torch.long).unsqueeze(0).to(device)

    with torch.no_grad():
        probs = model(input_ids, attention_mask)

    probs = probs.squeeze(0).cpu().tolist()  # 移除 batch 维度，并移到 CPU
    selected_sentences = []
    current_sentence = ""
    for i, token in enumerate(tokens):

        if token == '[CLS]' or token == '[SEP]':
          continue
        if token in ['，','。','？','！','；',';','!','?']:
          current_sentence += token
          if probs[i] > threshold:
            selected_sentences.append(current_sentence)
          current_sentence = ""
        else:
          current_sentence += token

    return "。".join(selected_sentences) + "。"

if __name__ == '__main__':
    # 加载模型和 tokenizer
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
    model = BertForExtractiveSummarization()
    model.load_state_dict(torch.load('extractive_summarizer.pth', map_location=device))  # 加载训练好的模型
    model.to(device)

    # 示例文本
    text = """
        你的长文本。
    """
    summary = summarize(text, model, tokenizer)
    print(f"Original Text:\n{text}\n")
    print(f"Summary:\n{summary}")

text_summarization_project/
├── data/
│   ├── train.txt     # 训练数据
│   ├── dev.txt       # 验证数据
│   └── test.txt      # 测试数据 (可选，用于最终评估)
├── data_utils.py    # 数据加载和预处理
├── model.py         # 模型定义
├── train.py         # 训练脚本
├── predict.py       # 预测/摘要脚本
└── extractive_summarizer.pth  # 保存的模型 (训练完成后)