Bert-mini2Bert-mini Summarization with 🤗EncoderDecoder Framework

This model is a warm-started BERT2BERT (mini) model fine-tuned on the CNN/Dailymail summarization dataset.

The model achieves a 16.51 ROUGE-2 score on CNN/Dailymail's test dataset.

For more details on how the model was fine-tuned, please refer to this notebook.

Results on test set 📝

Metric # Value
ROUGE-2 16.51

Model in Action 🚀

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

def generate_summary(text):
    # cut off at BERT max length 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 = "your text to be summarized here..."
generate_summary(text)

Created by Manuel Romero/@mrm8488 | LinkedIn

Made with in Spain

Downloads last month
200
Safetensors
Model size
23.4M params
Tensor type
I64
·
F32
·
Inference Examples
This model does not have enough activity to be deployed to Inference API (serverless) yet. Increase its social visibility and check back later, or deploy to Inference Endpoints (dedicated) instead.

Model tree for mrm8488/bert-mini2bert-mini-finetuned-cnn_daily_mail-summarization

Finetunes
2 models

Dataset used to train mrm8488/bert-mini2bert-mini-finetuned-cnn_daily_mail-summarization

Space using mrm8488/bert-mini2bert-mini-finetuned-cnn_daily_mail-summarization 1