Text Summarization with DistilBart Model

 

Text summarization represents a sophisticated evolution of text generation, requiring a deep understanding of content and context. With encoder-decoder transformer models like DistilBart, you can now create summaries that capture the essence of longer text while maintaining coherence and relevance.

In this tutorial, you’ll discover how to implement text summarization using DistilBart. You’ll learn through practical, executable examples, and by the end of this guide, you’ll understand both the theoretical foundations and hands-on implementation details. After completing this tutorial, you will know:

• How encoder-decoder transformers create context-aware summaries
• How to implement summarization using both pipeline and custom approaches
• How to build a production-ready summarization system with proper error handling
• How to optimize summarization parameters for different text types
• How to create a scalable summarization API service

Kick-start your project with my book NLP with Hugging Face Transformers. It provides self-study tutorials with working code.

Let’s get started.

 

Text Summarization with DistillBart Model
Photo by Aaron Burden. Some rights reserved.

Overview

This tutorial is in two parts; they are:

• Using DistilBart for Summarization
• Improving the Summarization Process

Using DistilBart for Summarization

Let’s start with a fundamental implementation that demonstrates the key concepts of text summarization with DistilBart:

import torch from transformers import AutoTokenizer, AutoModelForSeq2SeqLM   class TextSummarizer:     def __init__(self, model_name="sshleifer/distilbart-cnn-12-6"):         """Initialize the summarizer with a pre-trained model.           Args:             model_name (str): Name of the pre-trained model to use.         """         self.device = "cuda" if torch.cuda.is_available() else "cpu"         self.tokenizer = AutoTokenizer.from_pretrained(model_name)         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)         self.model.to(self.device)       def summarize(self, text, max_length=130, min_length=30, length_penalty=2.0,                   repetition_penalty=2.0, num_beams=4, early_stopping=True):         """Generate a summary for the given text.           Args:             text (str): The text to summarize             max_length (int): Maximum length of the summary             min_length (int): Minimum length of the summary             length_penalty (float): Penalty for longer summaries             repetition_penalty (float): Penalty for repeated tokens             num_beams (int): Number of beams for beam search             early_stopping (bool): Whether to stop when all beams are finished           Returns:             str: The generated summary         """         try:             # Tokenize the input text             inputs = self.tokenizer(text, max_length=1024, truncation=True,                                     padding="max_length", return_tensors="pt"                                     ).to(self.device)             # Generate summary             summary_ids = self.model.generate(                 inputs["input_ids"],                 attention_mask=inputs["attention_mask"],                 max_length=max_length,                 min_length=min_length,                 length_penalty=length_penalty,                 repetition_penalty=repetition_penalty,                 no_repeat_ngram_size=3,                 num_beams=num_beams,                 early_stopping=early_stopping             )             # Decode and return the summary             summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)             return summary         except Exception as e:             print(f"Error during summarization: {str(e)}")             return text   # Initialize the basic summarizer summarizer = TextSummarizer()   # Sample text to summarize sample_text = """ Artificial intelligence (AI) is a rapidly advancing field that focuses on creating intelligent systems capable of performing tasks that typically require human intelligence. These tasks include natural language processing, computer vision, speech recognition, and decision-making. With applications across healthcare, finance, education, and more, AI is transforming the way we interact with technology and solve complex problems. """   # Generate a summary summary = summarizer.summarize(sample_text) print("Basic Summary:\n", summary)   # Test with shorter text short_text = "AI is changing the world by automating tasks and providing insights from large datasets." short_summary = summarizer.summarize(short_text) print("Short Text Summary:\n", short_summary)

The output may be like this:

Basic Summary:   Artificial intelligence (AI) is transforming the way we interact with technology and   solve complex problems . The field focuses on creating intelligent systems capable of   performing tasks that typically require human intelligence . These tasks include natural   language processing, computer vision, speech recognition and decision-making . Short Text Summary:   AI is changing the world by automating tasks and providing insights from large datasets .   It's changing the way the world is changing to automatise tasks and provide insights .

Let’s break down to see how it works: The TextSummarizer class initializes with a pre-trained DistilBart model, specifically the CNN/DailyMail version that’s been optimized for summarization tasks. This is a powerful encoder-decoder transformer model.

Similar to many other models in the transformers library, you use the generate() method to pass on the input and retrieve the output. This is defined in the summarize() method of the TextSummarizer class. In particular, the generate() method takes the arguments max_length and min_length to control the length of the summary output. The arguments length_penalty controls the summary conciseness, and num_beams controls the beam search for better quality. The argument early_stopping is set to True to stop the generation when all beams are finished.

The model accepts only integer-encoded input, not text. The tokenizer object converts the input text into a sequence of subword tokens and then maps the tokens into integers. Its output will be a dictionary with the keys input_ids and attention_mask. The input_ids is the sequence of integers, and the attention_mask is a binary mask that indicates which tokens are real (1) and which are padding (0). Both the input_ids and attention_mask are required by the generate() method to understand your input correctly.

Finally, wrapping the generation code in a try-except block is a good practice to handle potential errors.

What is Beam Search?

The generate() method uses beam search to generate the summary. Recall that the model you used is an autoregressive model, which means it generates the output token by token. As the output token is sampled in each step, there are many possible output sequences. Unlike free-running text generation, a summary needs to be a complete paragraph, i.e., it should terminate with an end-of-sentence token.

Beam search is a technique used to generate summaries by exploring multiple paths in the output space until the end-of-sentence token is generated. Each step starts with 𝐾 partial sequences, and each will generate 𝐾 next token, so 𝐾2 sequences are considered. But only the top 𝐾 sequences of the highest score are kept for the next step.

The score of a sequence is computed based on multiple factors, including the probability of the tokens (i.e., the logit output of the model), the length of the sequence (length_penalty parameter), and the repetition of the sequence (repetition_penalty parameter). By adjusting these parameters, you can control the quality of the summary.

The no_repeat_ngram_size parameter is similar but not part of the score. Instead, it strictly forbids the repeated appearance of any n-grams in the output. You set it to 3 in the code above. Therefore, any consecutive three words will only appear once in the output. This is an additional check in the beam search algorithm.

Improving the Summarization Process

Using a pre-trained model from the Hugging Face model hub is as easy as above. You do not need to worry about the model architecture and the underlying implementation details. However, several issues may arise in real-world applications.

Quantization

Firstly, these models are large. You need to download a sizable model file before you can use it, and a lot of computation is required to generate a summary. You may find it too slow if you do not have a GPU. However, once the model is trained, you can approximate the floating point operations in the model with integer operations, which is much faster.

Converting a floating-point model to an integer is called quantization. Since the model is loaded as a PyTorch model, you can quantize it as follows:

... def __init__(self, model_name="sshleifer/distilbart-cnn-12-6", quantize=False):     """Initialize the advanced summarizer with additional features.       Args:         model_name (str): Name of the pre-trained model to use         quantize (bool): Whether to quantize the model for faster inference     """     self.device = "cuda" if torch.cuda.is_available() and not quantize else "cpu"     self.tokenizer = AutoTokenizer.from_pretrained(model_name)     self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)     if quantize:         self.model = torch.quantization.quantize_dynamic(self.model, dtype=torch.qint8)     self.model.to(self.device)

The if statement above converts the loaded model into a quantized version. The resulting model uses 8-bit integers, but not all operations are converted. PyTorch will convert the linear layers but likely not the attention layers, which will be kept as floating points.

Quantization may not work if you run your model on a GPU; it depends on the hardware and the version of PyTorch. You can try it out and see if it is supported in your environment.

Text Preprocessing and Batch Processing

As a summarizer, you know that some part of the input text is irrelevant, such as quoted URLs, extra whitespace, or special characters. You can remove them to avoid distracting the model.

Like many other machine learning models, the DistilBart model you used can accept not only a single string of text but a batch of texts. This is useful if you need to summarize multiple texts at once.

Let’s see how you can implement these:

... def preprocess_text(self, text: str) -> str:     """Clean and normalize input text.       Args:         text (str): Raw input text       Returns:         str: Cleaned and normalized text     """     # Remove extra whitespace     text = re.sub(r"\s+", " ", text.strip())     # Remove URLs     text = re.sub(r"^https?://[^\s/$.?#].[^\s]*$", "", text)     # Remove special characters but keep punctuation     text = re.sub(r"[^\w\s.,!?-]", "", text)       return text   def summarize_batch(self, texts: list[str], batch_size: int = 4, **kwargs) -> list[str]:     """Summarize multiple texts efficiently in batches.       Args:         texts (list[str]): List of input texts         batch_size (int): Number of texts to process at once         **kwargs: Additional arguments for summarization       Returns:         list[str]: List of generated summaries     """     summaries = []       for i in range(0, len(texts), batch_size):         # Create batch and process each text in the batch         batch = texts[i:i + batch_size]         processed_batch = [self.preprocess_text(text) for text in batch]           # Tokenize batch         inputs = self.tokenizer(processed_batch, max_length=1024, truncation=True,                                 padding=True, return_tensors="pt"                                 ).to(self.device)         # Generate summaries for batch         summary_ids = self.model.generate(             inputs["input_ids"],             attention_mask=inputs["attention_mask"],             **kwargs         )         # Decode summaries         summaries.extend([self.tokenizer.decode(ids, skip_special_tokens=True) for ids in summary_ids])     return summaries

The preprocess_text() method is the basic, and you can extend it. Here regular expressions are used to process the text to remove the unwanted elements in the input. In the summarize_batch() method, the input texts are tokenized and processed in batches. The summarize_batch() method accepts kwargs as additional arguments so that you can pass through the arguments such as length_penalty and num_beams to the generate() method.

Caching

Another issue when trying to use the model as a summarizer in production is employing caching. This means that when the same input is encountered again, the model will not reprocess it but retrieve the previous result. This not only saves time but also ensures consistency.

Surprisingly, implementing caching in a Python program is very straightforward. You can use the functools.lru_cache decorator to a function and the caching will be done automatically. Here, you can wrap the generation code in a decorator:

... @functools.lru_cache(maxsize=200) def cached_summarize(self, text: str, max_length: int = 130, min_length: int = 30,                         length_penalty: float = 2.0, repetition_penalty: float = 2.0,                         num_beams: int = 4, early_stopping: bool = True) -> str:     """Cached version of the summarization function."""     try:         # Tokenize the input text         inputs = self.tokenizer(text, max_length=1024, truncation=True,                                 padding="max_length", return_tensors="pt"                                 ).to(self.device)         # Generate summary         summary_ids = self.model.generate(             inputs["input_ids"],             attention_mask=inputs["attention_mask"],             max_length=max_length,             min_length=min_length,             length_penalty=length_penalty,             repetition_penalty=repetition_penalty,             num_beams=num_beams,             early_stopping=early_stopping,             no_repeat_ngram_size=3,         )         summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)         return summary     except Exception as e:         print(f"Error during summarization: {str(e)}")         return text

This is the magic provided by Python. As long as you call this function with the same arguments, the function will not be evaluated, but the previous result will be returned. The maxsize parameter controls the size of the cache and, hence, limits memory usage. Here, the cache will only remember the last 200 distinct results.

Long text

The model has a limit on the input size. It will refuse to run if you try to provide a very long input to the model. The max length for this model is 1024 tokens. You can check this limit from the config object or tokenizer:

import transformers   config = transformers.AutoConfig.from_pretrained("sshleifer/distilbart-cnn-12-6", trust_remote_code=True) print(config.max_position_embeddings)   tokenizer = transformers.AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6", trust_remote_code=True) print(tokenizer.model_max_length)

What if you want to summarize a text longer than 1024 tokens? You need to devise a strategy to split the text into manageable chunks so the model can handle it. Some options are:

• You can split the text into manageable chunks, summarizing each chunk separately and then combining the results.
• You consider the text as multiple paragraphs. Summarize each paragraph separately and then combine the results. The combined result will be a shorter text, which you can summarize again. Repeat this process until the input is short enough for the model to handle.

The first option is simpler, but you lost the context of the input. For example, if the key information is concentrated in the latter half of the input, you will generate a summary even for the first half. But since it is easier, let’s implement it:

... def split_long_text(self, text: str, max_tokens: int = 1024) -> list[str]:     """Split long text into chunks that fit within model's max token limit.       Args:         text (str): Input text         max_tokens (int): Maximum tokens per chunk       Returns:         list[str]: List of text chunks     """     # Tokenize the full text     tokens = self.tokenizer.tokenize(text)       # Split into chunks, then convert back to strings     chunks = [tokens[i:i+max_tokens] for i in range(0, len(tokens), max_tokens)]     return [self.tokenizer.convert_tokens_to_string(chunk) for chunk in chunks]   def summarize(self, text: str, max_length: int = 130, min_length: int = 30,                 length_penalty: float = 2.0, repetition_penalty: float = 2.0,                 num_beams: int = 4, early_stopping: bool = True) -> dict[str, str]:     """Generate a summary with advanced features.       Args:         text (str): The text to summarize         max_length (int): Maximum length of the summary         min_length (int): Minimum length of the summary         length_penalty (float): Penalty for longer summaries         repetition_penalty (float): Penalty for repeated tokens         num_beams (int): Number of beams for beam search         early_stopping (bool): Whether to stop when all beams are finished       Returns:         dict[str, str]: Dictionary containing original and summarized text     """     # Preprocess the text     cleaned_text = self.preprocess_text(text)       # Handle long texts     chunks = self.split_long_text(cleaned_text)     chunk_summaries = []     for chunk in chunks:         summary = self.cached_summarize(             chunk,             max_length=max_length // len(chunks),  # Adjust length for chunks             min_length=min_length // len(chunks),             length_penalty=length_penalty,             repetition_penalty=repetition_penalty,             num_beams=num_beams,             early_stopping=early_stopping         )         chunk_summaries.append(summary)       return {         "original_text": text,         "cleaned_text": cleaned_text,         "summary": " ".join(chunk_summaries)     }

The model is invoked using the cached_summarize() method you implemented earlier. Before that, you split the input into chunks; each is no longer than 1024 tokens, and process each chunk separately. To keep the output at a reasonable length, you adjust the max_length and min_length parameters accordingly so that the output after concatenation fits within the intended length.

Note that the model sees tokens, not words. Therefore, you split the input into chunks by counting the tokens and using the tokenizer to help in the splitting.

In full, below is the complete code:

import functools import pprint import re   import torch from transformers import AutoTokenizer, AutoModelForSeq2SeqLM     class AdvancedTextSummarizer:     def __init__(self, model_name="sshleifer/distilbart-cnn-12-6", quantize=False):         """Initialize the advanced summarizer with additional features.           Args:             model_name (str): Name of the pre-trained model to use             quantize (bool): Whether to quantize the model for faster inference         """         self.device = "cuda" if torch.cuda.is_available() and not quantize else "cpu"         self.tokenizer = AutoTokenizer.from_pretrained(model_name)         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)           if quantize:             self.model = torch.quantization.quantize_dynamic(self.model, dtype=torch.qint8)         self.model.to(self.device)       def preprocess_text(self, text: str) -> str:         """Clean and normalize input text.           Args:             text (str): Raw input text           Returns:             str: Cleaned and normalized text         """         # Remove extra whitespace         text = re.sub(r"\s+", " ", text.strip())         # Remove URLs         text = re.sub(r"^https?://[^\s/$.?#].[^\s]*$", "", text)         # Remove special characters but keep punctuation         text = re.sub(r"[^\w\s.,!?-]", "", text)           return text       def split_long_text(self, text: str, max_tokens: int = 1024) -> list[str]:         """Split long text into chunks that fit within model's max token limit.           Args:             text (str): Input text             max_tokens (int): Maximum tokens per chunk           Returns:             list[str]: List of text chunks         """         # Tokenize the full text         tokens = self.tokenizer.tokenize(text)           # Split into chunks, then convert back to strings         chunks = [tokens[i:i+max_tokens] for i in range(0, len(tokens), max_tokens)]         return [self.tokenizer.convert_tokens_to_string(chunk) for chunk in chunks]       @functools.lru_cache(maxsize=200)     def cached_summarize(self, text: str, max_length: int = 130, min_length: int = 30,                          length_penalty: float = 2.0, repetition_penalty: float = 2.0,                          num_beams: int = 4, early_stopping: bool = True) -> str:         """Cached version of the summarization function."""         try:             # Tokenize the input text             inputs = self.tokenizer(text, max_length=1024, truncation=True,                                     padding="max_length", return_tensors="pt"                                     ).to(self.device)             # Generate summary             summary_ids = self.model.generate(                 inputs["input_ids"],                 attention_mask=inputs["attention_mask"],                 max_length=max_length,                 min_length=min_length,                 length_penalty=length_penalty,                 repetition_penalty=repetition_penalty,                 num_beams=num_beams,                 early_stopping=early_stopping,                 no_repeat_ngram_size=3,   # Prevent repetition of phrases             )             summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)             return summary         except Exception as e:             print(f"Error during summarization: {str(e)}")             return text       def summarize_batch(self, texts: list[str], batch_size: int = 4, **kwargs) -> list[str]:         """Summarize multiple texts efficiently in batches.           Args:             texts (list[str]): List of input texts             batch_size (int): Number of texts to process at once             **kwargs: Additional arguments for summarization           Returns:             list[str]: List of generated summaries         """         summaries = []           for i in range(0, len(texts), batch_size):             # Create batch and process each text in the batch             batch = texts[i:i + batch_size]             processed_batch = [self.preprocess_text(text) for text in batch]               # Tokenize batch             inputs = self.tokenizer(processed_batch, max_length=1024, truncation=True,                                     padding=True, return_tensors="pt"                                     ).to(self.device)             # Generate summaries for batch             summary_ids = self.model.generate(                 inputs["input_ids"],                 attention_mask=inputs["attention_mask"],                 **kwargs             )             # Decode summaries             summaries.extend([self.tokenizer.decode(ids, skip_special_tokens=True) for ids in summary_ids])         return summaries       def summarize(self, text: str, max_length: int = 130, min_length: int = 30,                   length_penalty: float = 2.0, repetition_penalty: float = 2.0,                   num_beams: int = 4, early_stopping: bool = True) -> dict[str, str]:         """Generate a summary with advanced features.           Args:             text (str): The text to summarize             max_length (int): Maximum length of the summary             min_length (int): Minimum length of the summary             length_penalty (float): Penalty for longer summaries             repetition_penalty (float): Penalty for repeated tokens             num_beams (int): Number of beams for beam search             early_stopping (bool): Whether to stop when all beams are finished           Returns:             dict[str, str]: Dictionary containing original and summarized text         """         # Preprocess the text         cleaned_text = self.preprocess_text(text)           # Handle long texts         chunks = self.split_long_text(cleaned_text)         chunk_summaries = []         for chunk in chunks:             summary = self.cached_summarize(                 chunk,                 max_length=max_length // len(chunks),  # Adjust length for chunks                 min_length=min_length // len(chunks),                 length_penalty=length_penalty,                 repetition_penalty=repetition_penalty,                 num_beams=num_beams,                 early_stopping=early_stopping             )             chunk_summaries.append(summary)           return {             "original_text": text,             "cleaned_text": cleaned_text,             "summary": " ".join(chunk_summaries)         }   # Initialize the advanced summarizer with caching enabled and quantization adv_summarizer = AdvancedTextSummarizer(quantize=True)   # Sample text long_text = """ The development of artificial intelligence (AI) has significantly impacted various industries worldwide. From healthcare to finance, AI-powered applications have streamlined operations, improved accuracy, and unlocked new possibilities. In healthcare, AI assists in diagnostics, personalized treatment plans, and drug discovery. In finance, it aids in fraud detection, algorithmic trading, and customer service. Despite its benefits, AI raises concerns about data privacy, ethical implications, and job displacement. """   # Generate a summary with default settings adv_summary = adv_summarizer.summarize(long_text) print("Advanced Summary:") pprint.pprint(adv_summary)   # Batch summarization texts = [     "AI is revolutionizing healthcare with better diagnostics and personalized treatments.",     "Self-driving cars are powered by machine learning algorithms that continuously learn from traffic patterns.",     "Natural language processing helps computers understand and communicate with humans more effectively.",     "Climate change is being studied using AI models to predict future environmental patterns." ]   # Summarize multiple texts in a batch batch_summaries = adv_summarizer.summarize_batch(texts, batch_size=2) print("\nBatch Summaries:") for i, s in enumerate(batch_summaries, 1):     pprint.pprint(f"{i}. {s}")

Your output of the above code may be:

Advanced Summary: {'cleaned_text': 'The development of artificial intelligence AI has '                  'significantly impacted various industries worldwide. From '                  'healthcare to finance, AI-powered applications have '                  'streamlined operations, improved accuracy, and unlocked new '                  'possibilities. In healthcare, AI assists in diagnostics, '                  'personalized treatment plans, and drug discovery. In '                  'finance, it aids in fraud detection, algorithmic trading, '                  'and customer service. Despite its benefits, AI raises '                  'concerns about data privacy, ethical implications, and job '                  'displacement.', 'original_text': '\n'                   'The development of artificial intelligence (AI) has '                   'significantly impacted various industries worldwide.\n'                   'From healthcare to finance, AI-powered applications have '                   'streamlined operations, improved accuracy,\n'                   'and unlocked new possibilities. In healthcare, AI assists '                   'in diagnostics, personalized treatment plans,\n'                   'and drug discovery. In finance, it aids in fraud detection, '                   'algorithmic trading, and customer service.\n'                   'Despite its benefits, AI raises concerns about data '                   'privacy, ethical implications, and job displacement.\n', 'summary': ' From healthcare to healthcare, AI have streamlined operations '             'and improved accuracy . There is no doubt of the economy economy '             'economy in terms of economy economy . The economy economy has '             'been hit by economic stimulus stimulus stimulus. In addition to a '             'new economy at the hands of the public sector, the economy could '             'benefit from stimulus stimulus .'}   Batch Summaries: ('1.  AI is revolutionizing healthcare . AI is a software software software ' "that can help the world's healthcare workers .AI is a tool tool tool for " 'healthcare workers to learn more about the benefits of the company .AI has ' "been in the news for the first time, but it's not an easy work for the " "company, as well as the company's software software .") ('2.  Self-powered cars are powered powered by machine learning algorithms ' 'that learn from traffic traffic patterns . Self-learning algorithms are ' 'powered by algorithms algorithms that are constantly constantly changing ' 'traffic patterns in traffic patterns. Self-driving cars are driven by ' 'software software software that simulated traffic stops stops to stop ' 'traffic .') ('3.  "Natural language processing" helps the world around around the world . ' 'Natural gas processing technology is a tool tool tool for the most of of the ' "world's clocks clocks clocks clock clock clock clocks clocks . Natural " 'language processing technology can be used to understand the world of a ' 'power well-power .') ('4.  Climate change is a climate change topic topic of climate change change ' 'change . Climate change change changes changes changes in climate change . ' 'The climate change is being studied in the form form form of the world’s ' 'most-favavav global warming warming warming warmer warmer warmer ' 'temperatures .')

The code above uses pprint module to help print the output in a more readable format. Otherwise, all operations are trivial. You can try to tune the parameters to get a better summary.

Further Reading

Below are some resources that you may find useful:

• DistilBart model in Hugging Face
• Beam Search
• Foundations of NLP Explained Visually: Beam Search, How it Works

Summary

In this tutorial, you’ve learned how to summarize text using DistilBart. Particularly, you learned:

• Implement text summarization using DistilBart
• Build a custom summarizer with advanced features
• Process multiple texts efficiently
• Optimize summarization parameters

 

Want to Use Powerful Language Models in Your NLP Projects?