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. A seeded subset of 3,000 train dialogues is used for fine-tuning and 200 validation dialogues for the per-epoch loss. The test split has 500 unique dialogues, each with 3 human summaries. 200 of them are scored once, at the end.
  • Recipe: AdamW with lr 1e-3, weight decay 0.01 and linear decay to 0, batch size 8, 3 epochs (1,125 steps), gradient clipping at 1.0 and per-batch dynamic padding. Training is float32, or bf16 autocast on GPUs that support 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

Experiments

Weights pending. The full training run (flan-t5-base, 3,000 dialogues, 3 epochs) has not finished yet, so this repo does not contain the adapter or any results. The training notebook fills in this section and the table above from metrics.json when it exports the model.

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
-
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
(334)
this model

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

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