Summarization¶

The Task¶

Summarization is the task of summarizing text from a larger document/article into a short sentence/description. For example, taking a web article and describing the topic in a short sentence. This task is a subset of Sequence to Sequence tasks, which require the model to generate a variable length sequence given an input sequence. In our case the article would be our input sequence, and the short description/sentence would be the output sequence from the model.

Example¶

Let’s look at an example. We’ll use the XSUM dataset, which contains a train.csv and valid.csv. Each CSV file looks like this:

input,target
"The researchers have sequenced the genome of a strain of bacterium that causes the virulent infection...","A team of UK scientists hopes to shed light on the mysteries of bleeding canker, a disease that is threatening the nation's horse chestnut trees."
"Knight was shot in the leg by an unknown gunman at Miami's Shore Club where West was holding a pre-MTV Awards...",Hip hop star Kanye West is being sued by Death Row Records founder Suge Knight over a shooting at a beach party in August 2005.
...

In the above, the input column represents the long articles/documents, and the target is the short description used as the target. Once we’ve downloaded the data using download_data(), we create the SummarizationData. We select a pre-trained backbone to use for our SummarizationTask and finetune on the XSUM data. The backbone can be any Seq2Seq summarization model from HuggingFace/transformers.

Note

When changing the backbone, make sure you pass in the same backbone to the SummarizationData and the SummarizationTask!

Next, we use the trained SummarizationTask for inference. Finally, we save the model. Here’s the full example:

from flash import Trainer
from flash.core.data.utils import download_data
from flash.text import SummarizationData, SummarizationTask

# 1. Create the DataModule
download_data("https://pl-flash-data.s3.amazonaws.com/xsum.zip", "./data/")

datamodule = SummarizationData.from_csv(
    "input",
    "target",
    train_file="data/xsum/train.csv",
    val_file="data/xsum/valid.csv",
    batch_size=4,
)

# 2. Build the task
model = SummarizationTask()

# 3. Create the trainer and finetune the model
trainer = Trainer(max_epochs=3)
trainer.finetune(model, datamodule=datamodule, strategy="freeze")

# 4. Summarize some text!
datamodule = SummarizationData.from_lists(
    predict_data=[
        """
        Camilla bought a box of mangoes with a Brixton Â£10 note, introduced last year to try to keep the money of local
        people within the community.The couple were surrounded by shoppers as they walked along Electric Avenue.
        They came to Brixton to see work which has started to revitalise the borough.
        It was Charles' first visit to the area since 1996, when he was accompanied by the former
        South African president Nelson Mandela.Greengrocer Derek Chong, who has run a stall on Electric Avenue
        for 20 years, said Camilla had been ""nice and pleasant"" when she purchased the fruit.
        ""She asked me what was nice, what would I recommend, and I said we've got some nice mangoes.
        She asked me were they ripe and I said yes - they're from the Dominican Republic.""
        Mr Chong is one of 170 local retailers who accept the Brixton Pound.
        Customers exchange traditional pound coins for Brixton Pounds and then spend them at the market
        or in participating shops.
        During the visit, Prince Charles spent time talking to youth worker Marcus West, who works with children
        nearby on an estate off Coldharbour Lane. Mr West said:
        ""He's on the level, really down-to-earth. They were very cheery. The prince is a lovely man.""
        He added: ""I told him I was working with young kids and he said, 'Keep up all the good work.'""
        Prince Charles also visited the Railway Hotel, at the invitation of his charity The Prince's Regeneration Trust.
        The trust hopes to restore and refurbish the building,
        where once Jimi Hendrix and The Clash played, as a new community and business centre."
        """
    ],
    batch_size=4,
)
predictions = trainer.predict(model, datamodule=datamodule)
print(predictions)

# 5. Save the model!
trainer.save_checkpoint("summarization_model_xsum.pt")

To learn how to view the available backbones / heads for this task, see Backbones and Heads.

Flash Zero¶

The summarization task can be used directly from the command line with zero code using Flash Zero. You can run the above example with:

flash summarization

To view configuration options and options for running the summarization task with your own data, use:

flash summarization --help

Serving¶

The SummarizationTask is servable. This means you can call .serve to serve your Task. Here’s an example:

from flash.text import SummarizationTask

model = SummarizationTask.load_from_checkpoint(
    "https://flash-weights.s3.amazonaws.com/0.7.0/summarization_model_xsum.pt"
)
model.serve()

You can now perform inference from your client like this:

import requests

text = """
Camilla bought a box of mangoes with a Brixton Â£10 note, introduced last year to try to keep the money of local
people within the community.The couple were surrounded by shoppers as they walked along Electric Avenue.
They came to Brixton to see work which has started to revitalise the borough.
It was Charles' first visit to the area since 1996, when he was accompanied by the former
South African president Nelson Mandela.Greengrocer Derek Chong, who has run a stall on Electric Avenue
for 20 years, said Camilla had been ""nice and pleasant"" when she purchased the fruit.
""She asked me what was nice, what would I recommend, and I said we've got some nice mangoes.
She asked me were they ripe and I said yes - they're from the Dominican Republic.""
Mr Chong is one of 170 local retailers who accept the Brixton Pound.
Customers exchange traditional pound coins for Brixton Pounds and then spend them at the market
or in participating shops.
During the visit, Prince Charles spent time talking to youth worker Marcus West, who works with children
nearby on an estate off Coldharbour Lane. Mr West said:
""He's on the level, really down-to-earth. They were very cheery. The prince is a lovely man.""
He added: ""I told him I was working with young kids and he said, 'Keep up all the good work.'""
Prince Charles also visited the Railway Hotel, at the invitation of his charity The Prince's Regeneration Trust.
The trust hopes to restore and refurbish the building,
where once Jimi Hendrix and The Clash played, as a new community and business centre."
"""
body = {"session": "UUID", "payload": {"inputs": {"data": text}}}
resp = requests.post("http://127.0.0.1:8000/predict", json=body)

print(resp.json())

Accelerate Training & Inference with Torch ORT¶

Torch ORT converts your model into an optimized ONNX graph, speeding up training & inference when using NVIDIA or AMD GPUs. Enabling Torch ORT requires a single flag passed to the SummarizationTask once installed. See installation instructions here.

Note

Not all Transformer models are supported. See this table for supported models + branches containing fixes for certain models.

...

model = SummarizationTask(backbone="t5-large", num_classes=datamodule.num_classes, enable_ort=True)