Shortcuts

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."
        """,
        """
        "The problem is affecting people using the older versions of the PlayStation 3, called the ""Fat"" model. The
        problem isn't affecting the newer PS3 Slim systems that have been on sale since September last year. Sony have
        also said they are aiming to have the problem fixed shortly but is advising some users to avoid using their
        console for the time being.""We hope to resolve this problem within the next 24 hours,"" a statement reads.
        ""In the meantime, if you have a model other than the new slim PS3, we advise that you do not use your PS3
        system, as doing so may result in errors in some functionality, such as recording obtained trophies, and not
        being able to restore certain data.""We believe we have identified that this problem is being caused by a bug
        in the clock functionality incorporated in the system.""The PlayStation Network is used by millions of people
        around the world.It allows users to play their friends at games like Fifa over the internet and also do things
        like download software or visit online stores.",Sony has told owners of older models of its PlayStation 3
        console to stop using the machine because of a problem with the PlayStation Network.
        """,
    ],
    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)
Read the Docs v: latest
Versions
latest
stable
0.8.2
0.8.1.post0
0.8.1
0.8.0
0.7.5
0.7.4
0.7.3
0.7.2
0.7.1
0.7.0
0.6.0
0.5.2
0.5.1
0.5.0
0.4.0
0.3.2
0.3.1
0.3.0
0.2.3
0.2.2
0.2.1
0.2.0
0.1.0post1
Downloads
html
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.