FLAN-T5 Dialogue Summarizer + RAG

A LoRA adapter that fine-tunes google/flan-t5-base to write short third-person summaries of everyday conversations, trained on DialogSum. The same model also answers questions about a fictional electronics shop, "Nova Gadgets". A MiniLM retriever finds the relevant passages of a 41-document knowledge base (assets/kb.json), and FLAN-T5 answers from those passages only (retrieval-augmented generation).

Model

  • Base: google/flan-t5-base, an instruction-tuned T5 encoder-decoder (247,577,856 parameters, float32).
  • Adapter: LoRA with rank 16, alpha 32 and dropout 0.05 on the attention query and value projections (q, v) of every encoder and decoder layer. That is 1,769,472 trainable parameters (0.71 % of the base). The base weights stay frozen.
  • Files: adapter_config.json + adapter_model.safetensors (the adapter, about 7 MB), the tokenizer, config.json (base model id, generation settings, prompt templates, library versions) and assets/kb.json (the RAG knowledge base).
  • Loading (model.load): the base model is downloaded from the Hub, the adapter is merged into it on the CPU (merge_and_unload), and the merged model is then moved to the device. Inference therefore costs exactly what the base model costs.
  • Prompt: Summarize the following conversation.\n\n{dialogue}\n\nSummary: , the same template in training, evaluation and serving. Input is truncated to 512 tokens. Decoding is greedy with at most 96 new tokens.
  • RAG: sentence-transformers/all-MiniLM-L6-v2 embeds the question and the 41 documents on the CPU. The top 3 documents by cosine similarity go into the prompt Answer the question using only the context below. ..., which the same merged model answers.

Outputs:

  • predict(dialogue) -> str: the summary.
  • rag(question) -> {"answer": str, "sources": [{"id", "title", "text", "score"}, ...]}.

Usage

from huggingface_hub import snapshot_download
import sys
path = snapshot_download("shalev396/flan-t5-dialogue-summarizer")
sys.path.insert(0, path)
import model
predictor = model.load(path, device="cpu")   # or "cuda"
print(predictor.predict("#Person1#: Is the 7:15 train on time?\n#Person2#: No, it's 20 minutes late.\n"
                        "#Person1#: Then I'll grab a coffee first."))
print(predictor.rag("How long does standard shipping take?"))

model.load_base(path) gives the same base model without the adapter, for comparisons.

With PEFT only: PeftModel.from_pretrained(AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base"), "shalev396/flan-t5-dialogue-summarizer"). Use the prompt above.

  • Space / free API: shalev396/flan-t5-dialogue-summarizer, endpoints /predict (dialogue -> summary) and /rag (question -> answer + sources).
  • Inference Endpoint: handler.py. Send {"inputs": "<dialogue>"} to get {"summary": ...}, or {"inputs": "<question>", "parameters": {"task": "rag"}} to get {"answer": ..., "sources": [...]}.

Training

  • Data: DialogSum has 13,460 dialogues with human summaries. All 12,460 train dialogues are used for fine-tuning and all 500 validation dialogues for the per-epoch loss. The test split has 500 unique dialogues, each with 3 human summaries, all scored once, at the end.
  • Recipe: AdamW with lr 1e-3, weight decay 0.01 and linear decay to 0, batch size 8, up to 6 epochs of 1,558 steps, gradient clipping at 1.0 and per-batch dynamic padding. The adapter of the epoch with the lowest validation loss is kept (epoch 5 of 6), and training stops early after 2 epochs without a new best. Training is float32, or bf16 autocast on GPUs that support bf16 (this run: Google Colab, NVIDIA A100, bf16). fp16 is never used, because T5 overflows in fp16.
  • Variants: the zero-shot, one-shot and few-shot (k=2) base model, and the LoRA fine-tune, all scored on the same test dialogues with greedy decoding.
  • Full code: training/ · Colab

Evaluation

metric (test) value
rouge1 0.4500
rouge2 0.1904
rougeL 0.3660

Experiments

ROUGE F-measure x100 on 500 DialogSum test dialogues, each scored against its 3 human summaries (mean). Greedy decoding for every variant. The adapter adds +16.34 ROUGE-L over the zero-shot base model.

variant ROUGE-1 ROUGE-2 ROUGE-L avg. words
base zero-shot 23.53 7.21 20.27 14.2
base one-shot 24.13 7.11 20.63 16.3
base few-shot (k=2) 24.31 7.01 20.75 15.9
LoRA fine-tuned (deployed) 45.00 19.04 36.60 23.4

RAG retriever (MiniLM, 18 held-out store questions): hit@1 0.94 · hit@3 1.00 · MRR 0.97.

Run: trained 2026-09-27 with training/ on Google Colab (NVIDIA A100) · bf16 autocast · train 1788 s (0.19 s/step) · eval 268 s.

Loss curve ROUGE comparison ROUGE-L per test dialogue

Limitations

  • English chit-chat only. DialogSum is written, two-speaker English everyday dialogue (#Person1#: / #Person2#: turns). Meetings, transcripts with many speakers, and other languages are out of domain.
  • Summaries can be wrong. A 250M-parameter model sometimes swaps who said what, drops key facts or invents details. Check anything that matters against the conversation.
  • ROUGE measures word overlap, not factual correctness. The test set has only 200 dialogues, so small differences between variants are within noise.
  • Long inputs are cut at 512 tokens. That affects about 2 % of DialogSum train prompts (the end of the conversation is lost).
  • RAG is a toy. The knowledge base describes a fictional store, and the answer generator is the summarization fine-tune, not a QA model. Answers can be copied sentences or miss part of a question. The retriever always returns 3 documents, even for off-topic questions.
  • Educational portfolio project.
Downloads last month
45
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for shalev396/flan-t5-dialogue-summarizer

Adapter
(335)
this model

Dataset used to train shalev396/flan-t5-dialogue-summarizer

Space using shalev396/flan-t5-dialogue-summarizer 1

Evaluation results