mBART-50 Dialogue Summarization
This model is a fine-tuned version of facebook/mbart-large-50-many-to-many-mmt for English dialogue summarization.
It takes an English multi-speaker dialogue as input and generates an English abstractive summary.
Model Details
- Base model:
facebook/mbart-large-50-many-to-many-mmt - Task: Dialogue Summarization
- Input language: English
- Output language: English
- Framework: Hugging Face Transformers
- Model type: Seq2Seq Transformer
- Fine-tuned by: Yunwoo Lee
Intended Use
This model is intended for:
- English dialogue summarization
- Research on abstractive summarization
- Educational and experimental use with sequence-to-sequence models
- Benchmarking or reproducing dialogue summarization pipelines
This model is not intended for high-stakes or production use without additional evaluation.
Training Data
The model was fine-tuned on a JSON-formatted dialogue summarization dataset with the following fields in the raw data:
dialoguesummarysummary_desummary_zh
For this training run, only the following columns were used:
dialogue→ source textsummary→ target summary
After cleaning and filtering, the dataset sizes were:
- Train: 14,730
- Validation: 818
- Test: 819
Preprocessing
The following preprocessing steps were applied:
- Removed placeholder tokens such as
<file_gif>and similar markers - Normalized whitespace
- Preserved useful speaker information in the dialogue
- Dropped empty or invalid rows after cleaning
Tokenization Setup
- Source max length: 768
- Target max length: 64
- Source language code:
en_XX - Target language code:
en_XX
Training Procedure
The model was trained with the following configuration:
- Epochs: 5
- Learning rate:
2e-5 - Optimizer:
Adafactor - Per-device train batch size: 4
- Gradient accumulation steps: 6
- Effective batch size: 24
- Per-device eval batch size: 4
- Beam size: 5
- Mixed precision:
bf16 - Best model selection metric:
rougeLsum
Hardware
Training was run on:
- GPU: NVIDIA A100-SXM4-40GB
- Environment: Google Colab
- PyTorch: 2.10.0+cu128
Evaluation
Evaluation was performed using ROUGE.
Validation Results
| Metric | Score |
|---|---|
| ROUGE-1 | 50.9603 |
| ROUGE-2 | 26.7346 |
| ROUGE-L | 42.1370 |
| ROUGE-Lsum | 46.6278 |
| Loss | 1.4443 |
Test Results
| Metric | Score |
|---|---|
| ROUGE-1 | 48.9213 |
| ROUGE-2 | 24.1386 |
| ROUGE-L | 40.3879 |
| ROUGE-Lsum | 44.7998 |
| Loss | 1.4613 |
Recommended Generation Settings
The model was evaluated using beam search.
For best results, use generation settings close to the evaluation setup:
num_beams=5max_length=64forced_bos_token_id=en_XX
How to Use
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
model_name = "yunu919/mbart-large-dialogue-summarization"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
dialogue = """Hannah: Hey, do you have Betty's number?
Amanda: Lemme check
Amanda: Sorry, can't find it.
Amanda: Ask Larry"""
inputs = tokenizer(
dialogue,
return_tensors="pt",
truncation=True,
max_length=768,
)
outputs = model.generate(
**inputs,
num_beams=5,
max_length=64,
forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"],
)
summary = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
print(summary)