pytorch实现文本摘要

复制代码

import numpy as np

from modelscope.hub.snapshot_download import snapshot_download
from transformers import BertTokenizer, BertModel
import torch

# 下载模型到本地目录
model_dir = snapshot_download('tiansz/bert-base-chinese', cache_dir='./bert-base-chinese')
print(f"模型已下载到: {model_dir}")

# 本地模型路径
model_path = model_dir  # 使用下载的模型路径

# 从本地加载分词器和模型
tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertModel.from_pretrained(model_path)


def get_sentence_embeddings(text):
    """
    获取输入文本的句子级别嵌入（BERT模型的输出）用于摘要任务
    """
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
    # 获取最后一层的[CLS] token的嵌入向量作为句子的表示
    return outputs.pooler_output.detach().numpy()


def summarize(text, num_sentences=3):
    """
    使用抽取式摘要从输入文本中提取最重要的句子
    """
    sentences = text.split("。")  # 以句号分割句子
    sentence_embeddings = []

    for sentence in sentences:
        embedding = get_sentence_embeddings(sentence)
        sentence_embeddings.append(embedding)

    # 使用句子得分来排序
    sentence_scores = np.array([embedding[0][0] for embedding in sentence_embeddings])
    ranked_sentences = [sentences[i] for i in sentence_scores.argsort()[-num_sentences:][::-1]]

    # 返回排名靠前的句子
    return "。".join(ranked_sentences)


# 示例中文文本
text = """
机器学习（Machine Learning，简称 ML）是人工智能（AI）领域的一个分支，重点研究开发能够使计算机在没有明确编程指令的情况下进行任务的算法和统计模型。机器学习通过从数据中学习，识别模式，并根据少量人为干预做出决策。 
近年来，深度学习（Deep Learning）作为机器学习的一个子集，取得了显著的进展，尤其是在计算机视觉、自然语言处理和语音识别等领域。 
深度学习技术使用多层神经网络结构，能够通过处理大量的数据来自动提取特征，从而提高机器学习的性能。 
目前，人工智能已经在多个行业中得到广泛应用，包括医疗健康、金融、自动驾驶等领域。
"""

# 获取摘要
summary = summarize(text)
print("原文：\n", text)
print("\n摘要：\n", summary)

～困了🥱，通宵写了一个晚上

结果：

复制代码

Downloading Model to directory: ./bert-base-chinese/tiansz/bert-base-chinese
模型已下载到: ./bert-base-chinese/tiansz/bert-base-chinese
Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.
原文：
 
机器学习（Machine Learning，简称 ML）是人工智能（AI）领域的一个分支，重点研究开发能够使计算机在没有明确编程指令的情况下进行任务的算法和统计模型。机器学习通过从数据中学习，识别模式，并根据少量人为干预做出决策。 
近年来，深度学习（Deep Learning）作为机器学习的一个子集，取得了显著的进展，尤其是在计算机视觉、自然语言处理和语音识别等领域。 
深度学习技术使用多层神经网络结构，能够通过处理大量的数据来自动提取特征，从而提高机器学习的性能。 
目前，人工智能已经在多个行业中得到广泛应用，包括医疗健康、金融、自动驾驶等领域。


摘要：
  
目前，人工智能已经在多个行业中得到广泛应用，包括医疗健康、金融、自动驾驶等领域。机器学习通过从数据中学习，识别模式，并根据少量人为干预做出决策。 
深度学习技术使用多层神经网络结构，能够通过处理大量的数据来自动提取特征，从而提高机器学习的性能