语言: 英语
许可证: Apache-2.0
数据集:

• cnn_dailymail
  标签:
• 摘要生成

基于🤗编码器-解码器框架的Bert-small2Bert-small摘要模型

本模型是一个经过热启动的BERT2BERT（小型版）模型，并在CNN/Dailymail摘要数据集上进行了微调。

该模型在CNN/Dailymail测试集上取得了17.37的ROUGE-2分数。

关于模型微调的更多细节，请参考此笔记本。

测试集结果 📝

模型实战 🚀

from transformers import BertTokenizerFast, EncoderDecoderModel  
import torch  
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
tokenizer = BertTokenizerFast.from_pretrained('mrm8488/bert-small2bert-small-finetuned-cnn_daily_mail-summarization')  
model = EncoderDecoderModel.from_pretrained('mrm8488/bert-small2bert-small-finetuned-cnn_daily_mail-summarization').to(device)  

def generate_summary(text):  
    # 在BERT最大长度512处截断  
    inputs = tokenizer([text], padding="max_length", truncation=True, max_length=512, return_tensors="pt")  
    input_ids = inputs.input_ids.to(device)  
    attention_mask = inputs.attention_mask.to(device)  

    output = model.generate(input_ids, attention_mask=attention_mask)  

    return tokenizer.decode(output[0], skip_special_tokens=True)  

text = "待摘要的文本放在这里..."  
generate_summary(text)  

由Manuel Romero/@mrm8488创建 | LinkedIn

制作于西班牙，充满❤️