We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Task doesn't seem to converge: Summurization, Translation
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

On a machine with 1 GPU.

```
FLASH_TESTING=1 python flash_examples/finetuning/summarization.py
```

```
FLASH_TESTING=1 python flash_examples/finetuning/translation.py
```

Steps to reproduce the behavior:

1. Go to '...'
2. Run '....'
3. Scroll down to '....'
4. See error

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment

 - PyTorch Version (e.g., 1.0):
 - OS (e.g., Linux):
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version:
 - GPU models and configuration:
 - Any other relevant information:

### Additional context

<!-- Add any other context about the problem here. -->


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flash/text/seq2seq/summarization/data.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from flash.text.seq2seq.core.data import Seq2SeqData, Seq2SeqPostprocess, Seq2SeqPreprocess
15 
16 
17 class SummarizationData(Seq2SeqData):
18 
19     preprocess_cls = Seq2SeqPreprocess
20     postprocess_cls = Seq2SeqPostprocess
21
```

Path: `flash/text/seq2seq/summarization/metric.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from typing import Dict, List, Tuple
15 
16 import numpy as np
17 from torch import tensor
18 from torchmetrics import Metric
19 
20 from flash.core.utilities.imports import _TEXT_AVAILABLE
21 from flash.text.seq2seq.summarization.utils import add_newline_to_end_of_each_sentence
22 
23 if _TEXT_AVAILABLE:
24     from rouge_score import rouge_scorer, scoring
25     from rouge_score.scoring import AggregateScore, BootstrapAggregator, Score
26 else:
27     AggregateScore, Score, BootstrapAggregator = None, None, object
28 
29 
30 class RougeMetric(Metric):
31     """
32     Metric used for automatic summarization. https://www.aclweb.org/anthology/W04-1013/
33 
34     Example:
35 
36         >>> target = "Is your name John".split()
37         >>> preds = "My name is John".split()
38         >>> rouge = RougeMetric()   # doctest: +SKIP
39         >>> from pprint import pprint
40         >>> pprint(rouge(preds, target))  # doctest: +NORMALIZE_WHITESPACE +SKIP
41         {'rouge1_fmeasure': 0.25,
42          'rouge1_precision': 0.25,
43          'rouge1_recall': 0.25,
44          'rouge2_fmeasure': 0.0,
45          'rouge2_precision': 0.0,
46          'rouge2_recall': 0.0,
47          'rougeL_fmeasure': 0.25,
48          'rougeL_precision': 0.25,
49          'rougeL_recall': 0.25,
50          'rougeLsum_fmeasure': 0.25,
51          'rougeLsum_precision': 0.25,
52          'rougeLsum_recall': 0.25}
53     """
54 
55     def __init__(
56         self,
57         rouge_newline_sep: bool = False,
58         use_stemmer: bool = False,
59         rouge_keys: Tuple[str] = ("rouge1", "rouge2", "rougeL", "rougeLsum"),
60     ):
61         super().__init__()
62         if not _TEXT_AVAILABLE:
63             raise ModuleNotFoundError("Please, pip install 'lightning-flash[text]'")
64 
65         self.rouge_newline_sep = rouge_newline_sep
66         self.rouge_keys = rouge_keys
67         self.use_stemmer = use_stemmer
68         self.aggregator = RougeBatchAggregator()
69         self.scorer = rouge_scorer.RougeScorer(rouge_keys, use_stemmer=self.use_stemmer)
70 
71         for key in rouge_keys:
72             self.add_state(key, [])
73 
74     def update(self, pred_lns: List[str], tgt_lns: List[str]):
75         for pred, tgt in zip(pred_lns, tgt_lns):
76             # rougeLsum expects "\n" separated sentences within a summary
77             if self.rouge_newline_sep:
78                 pred = add_newline_to_end_of_each_sentence(pred)
79                 tgt = add_newline_to_end_of_each_sentence(tgt)
80             results = self.scorer.score(pred, tgt)
81             for key, score in results.items():
82                 score = tensor([score.precision, score.recall, score.fmeasure])
83                 getattr(self, key).append(score)
84 
85     def compute(self) -> Dict[str, float]:
86         scores = {key: getattr(self, key) for key in self.rouge_keys}
87         self.aggregator.add_scores(scores)
88         result = self.aggregator.aggregate()
89         return format_rouge_results(result)
90 
91     def __hash__(self):
92         # override to hash list objects.
93         # this is a bug in the upstream pytorch release.
94         hash_vals = [self.__class__.__name__]
95 
96         for key in self._defaults.keys():
97             value = getattr(self, key)
98             if isinstance(value, list):
99                 value = tuple(value)
100             hash_vals.append(value)
101 
102         return hash(tuple(hash_vals))
103 
104 
105 class RougeBatchAggregator(BootstrapAggregator):
106     """
107     Aggregates rouge scores and provides confidence intervals.
108     """
109 
110     def aggregate(self):
111         """
112         Override function to wrap the final results in `Score` objects.
113         This is due to the scores being replaced with a list of torch tensors.
114         """
115         result = {}
116         for score_type, scores in self._scores.items():
117             # Stack scores into a 2-d matrix of (sample, measure).
118             score_matrix = np.vstack(tuple(scores))
119             # Percentiles are returned as (interval, measure).
120             percentiles = self._bootstrap_resample(score_matrix)
121             # Extract the three intervals (low, mid, high).
122             intervals = tuple((Score(*percentiles[j, :]) for j in range(3)))
123             result[score_type] = AggregateScore(low=intervals[0], mid=intervals[1], high=intervals[2])
124         return result
125 
126     def add_scores(self, scores):
127         self._scores = scores
128 
129 
130 def format_rouge_results(result: Dict[str, AggregateScore], decimal_places: int = 4) -> Dict[str, float]:
131     flattened_result = {}
132     for rouge_key, rouge_aggregate_score in result.items():
133         for stat in ["precision", "recall", "fmeasure"]:
134             mid = rouge_aggregate_score.mid
135             score = round(getattr(mid, stat), decimal_places)
136             flattened_result[f"{rouge_key}_{stat}"] = score
137     return flattened_result
138
```

Path: `flash/text/seq2seq/summarization/model.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union
15 
16 import pytorch_lightning as pl
17 import torch
18 
19 from flash.text.seq2seq.core.model import Seq2SeqTask
20 from flash.text.seq2seq.summarization.metric import RougeMetric
21 
22 
23 class SummarizationTask(Seq2SeqTask):
24     """Task for Seq2Seq Summarization.
25 
26     Args:
27         backbone: backbone model to use for the task.
28         loss_fn: Loss function for training.
29         optimizer: Optimizer to use for training, defaults to `torch.optim.Adam`.
30         metrics: Metrics to compute for training and evaluation.
31         learning_rate: Learning rate to use for training, defaults to `3e-4`
32         val_target_max_length: Maximum length of targets in validation. Defaults to `128`
33         num_beams: Number of beams to use in validation when generating predictions. Defaults to `4`
34         use_stemmer: Whether Porter stemmer should be used to strip word suffixes to improve matching.
35         rouge_newline_sep: Add a new line at the beginning of each sentence in Rouge Metric calculation.
36     """
37 
38     def __init__(
39         self,
40         backbone: str = "sshleifer/tiny-mbart",
41         loss_fn: Optional[Union[Callable, Mapping, Sequence]] = None,
42         optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
43         metrics: Union[pl.metrics.Metric, Mapping, Sequence, None] = None,
44         learning_rate: float = 5e-5,
45         val_target_max_length: Optional[int] = None,
46         num_beams: Optional[int] = 4,
47         use_stemmer: bool = True,
48         rouge_newline_sep: bool = True
49     ):
50         self.save_hyperparameters()
51         super().__init__(
52             backbone=backbone,
53             loss_fn=loss_fn,
54             optimizer=optimizer,
55             metrics=metrics,
56             learning_rate=learning_rate,
57             val_target_max_length=val_target_max_length,
58             num_beams=num_beams
59         )
60         self.rouge = RougeMetric(
61             rouge_newline_sep=rouge_newline_sep,
62             use_stemmer=use_stemmer,
63         )
64 
65     @property
66     def task(self) -> str:
67         return "summarization"
68 
69     def compute_metrics(self, generated_tokens: torch.Tensor, batch: Dict, prefix: str) -> None:
70         tgt_lns = self.tokenize_labels(batch["labels"])
71         result = self.rouge(self._postprocess.uncollate(generated_tokens), tgt_lns)
72         self.log_dict(result, on_step=False, on_epoch=True)
73 
74     def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):
75         """
76         This function is used only for debugging usage with CI
77         """
78         assert history[-1]["val_f1"] > 0.45
79
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/flash/text/seq2seq/summarization/data.py b/flash/text/seq2seq/summarization/data.py
--- a/flash/text/seq2seq/summarization/data.py
+++ b/flash/text/seq2seq/summarization/data.py
@@ -11,10 +11,37 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Callable, Dict, Optional, Union
+
 from flash.text.seq2seq.core.data import Seq2SeqData, Seq2SeqPostprocess, Seq2SeqPreprocess
 
 
+class SummarizationPreprocess(Seq2SeqPreprocess):
+
+    def __init__(
+        self,
+        train_transform: Optional[Dict[str, Callable]] = None,
+        val_transform: Optional[Dict[str, Callable]] = None,
+        test_transform: Optional[Dict[str, Callable]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        backbone: str = "sshleifer/distilbart-xsum-1-1",
+        max_source_length: int = 128,
+        max_target_length: int = 128,
+        padding: Union[str, bool] = 'max_length'
+    ):
+        super().__init__(
+            train_transform=train_transform,
+            val_transform=val_transform,
+            test_transform=test_transform,
+            predict_transform=predict_transform,
+            backbone=backbone,
+            max_source_length=max_source_length,
+            max_target_length=max_target_length,
+            padding=padding,
+        )
+
+
 class SummarizationData(Seq2SeqData):
 
-    preprocess_cls = Seq2SeqPreprocess
+    preprocess_cls = SummarizationPreprocess
     postprocess_cls = Seq2SeqPostprocess
diff --git a/flash/text/seq2seq/summarization/metric.py b/flash/text/seq2seq/summarization/metric.py
--- a/flash/text/seq2seq/summarization/metric.py
+++ b/flash/text/seq2seq/summarization/metric.py
@@ -21,7 +21,7 @@
 from flash.text.seq2seq.summarization.utils import add_newline_to_end_of_each_sentence
 
 if _TEXT_AVAILABLE:
-    from rouge_score import rouge_scorer, scoring
+    from rouge_score import rouge_scorer
     from rouge_score.scoring import AggregateScore, BootstrapAggregator, Score
 else:
     AggregateScore, Score, BootstrapAggregator = None, None, object
diff --git a/flash/text/seq2seq/summarization/model.py b/flash/text/seq2seq/summarization/model.py
--- a/flash/text/seq2seq/summarization/model.py
+++ b/flash/text/seq2seq/summarization/model.py
@@ -37,11 +37,11 @@
 
     def __init__(
         self,
-        backbone: str = "sshleifer/tiny-mbart",
+        backbone: str = "sshleifer/distilbart-xsum-1-1",
         loss_fn: Optional[Union[Callable, Mapping, Sequence]] = None,
         optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
         metrics: Union[pl.metrics.Metric, Mapping, Sequence, None] = None,
-        learning_rate: float = 5e-5,
+        learning_rate: float = 1e-5,
         val_target_max_length: Optional[int] = None,
         num_beams: Optional[int] = 4,
         use_stemmer: bool = True,
@@ -69,10 +69,10 @@
     def compute_metrics(self, generated_tokens: torch.Tensor, batch: Dict, prefix: str) -> None:
         tgt_lns = self.tokenize_labels(batch["labels"])
         result = self.rouge(self._postprocess.uncollate(generated_tokens), tgt_lns)
-        self.log_dict(result, on_step=False, on_epoch=True)
+        self.log_dict(result, on_step=False, on_epoch=True, prog_bar=True)
 
     def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):
         """
         This function is used only for debugging usage with CI
         """
-        assert history[-1]["val_f1"] > 0.45
+        assert history[-1]["rouge1_recall"] > 0.2

{"golden_diff": "diff --git a/flash/text/seq2seq/summarization/data.py b/flash/text/seq2seq/summarization/data.py\n--- a/flash/text/seq2seq/summarization/data.py\n+++ b/flash/text/seq2seq/summarization/data.py\n@@ -11,10 +11,37 @@\n # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n # See the License for the specific language governing permissions and\n # limitations under the License.\n+from typing import Callable, Dict, Optional, Union\n+\n from flash.text.seq2seq.core.data import Seq2SeqData, Seq2SeqPostprocess, Seq2SeqPreprocess\n \n \n+class SummarizationPreprocess(Seq2SeqPreprocess):\n+\n+    def __init__(\n+        self,\n+        train_transform: Optional[Dict[str, Callable]] = None,\n+        val_transform: Optional[Dict[str, Callable]] = None,\n+        test_transform: Optional[Dict[str, Callable]] = None,\n+        predict_transform: Optional[Dict[str, Callable]] = None,\n+        backbone: str = \"sshleifer/distilbart-xsum-1-1\",\n+        max_source_length: int = 128,\n+        max_target_length: int = 128,\n+        padding: Union[str, bool] = 'max_length'\n+    ):\n+        super().__init__(\n+            train_transform=train_transform,\n+            val_transform=val_transform,\n+            test_transform=test_transform,\n+            predict_transform=predict_transform,\n+            backbone=backbone,\n+            max_source_length=max_source_length,\n+            max_target_length=max_target_length,\n+            padding=padding,\n+        )\n+\n+\n class SummarizationData(Seq2SeqData):\n \n-    preprocess_cls = Seq2SeqPreprocess\n+    preprocess_cls = SummarizationPreprocess\n     postprocess_cls = Seq2SeqPostprocess\ndiff --git a/flash/text/seq2seq/summarization/metric.py b/flash/text/seq2seq/summarization/metric.py\n--- a/flash/text/seq2seq/summarization/metric.py\n+++ b/flash/text/seq2seq/summarization/metric.py\n@@ -21,7 +21,7 @@\n from flash.text.seq2seq.summarization.utils import add_newline_to_end_of_each_sentence\n \n if _TEXT_AVAILABLE:\n-    from rouge_score import rouge_scorer, scoring\n+    from rouge_score import rouge_scorer\n     from rouge_score.scoring import AggregateScore, BootstrapAggregator, Score\n else:\n     AggregateScore, Score, BootstrapAggregator = None, None, object\ndiff --git a/flash/text/seq2seq/summarization/model.py b/flash/text/seq2seq/summarization/model.py\n--- a/flash/text/seq2seq/summarization/model.py\n+++ b/flash/text/seq2seq/summarization/model.py\n@@ -37,11 +37,11 @@\n \n     def __init__(\n         self,\n-        backbone: str = \"sshleifer/tiny-mbart\",\n+        backbone: str = \"sshleifer/distilbart-xsum-1-1\",\n         loss_fn: Optional[Union[Callable, Mapping, Sequence]] = None,\n         optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,\n         metrics: Union[pl.metrics.Metric, Mapping, Sequence, None] = None,\n-        learning_rate: float = 5e-5,\n+        learning_rate: float = 1e-5,\n         val_target_max_length: Optional[int] = None,\n         num_beams: Optional[int] = 4,\n         use_stemmer: bool = True,\n@@ -69,10 +69,10 @@\n     def compute_metrics(self, generated_tokens: torch.Tensor, batch: Dict, prefix: str) -> None:\n         tgt_lns = self.tokenize_labels(batch[\"labels\"])\n         result = self.rouge(self._postprocess.uncollate(generated_tokens), tgt_lns)\n-        self.log_dict(result, on_step=False, on_epoch=True)\n+        self.log_dict(result, on_step=False, on_epoch=True, prog_bar=True)\n \n     def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):\n         \"\"\"\n         This function is used only for debugging usage with CI\n         \"\"\"\n-        assert history[-1][\"val_f1\"] > 0.45\n+        assert history[-1][\"rouge1_recall\"] > 0.2\n", "issue": "Task doesn't seem to converge: Summurization, Translation\n## \ud83d\udc1b Bug\r\n\r\n<!-- A clear and concise description of what the bug is. -->\r\n\r\n### To Reproduce\r\n\r\nOn a machine with 1 GPU.\r\n\r\n```\r\nFLASH_TESTING=1 python flash_examples/finetuning/summarization.py\r\n```\r\n\r\n```\r\nFLASH_TESTING=1 python flash_examples/finetuning/translation.py\r\n```\r\n\r\nSteps to reproduce the behavior:\r\n\r\n1. Go to '...'\r\n2. Run '....'\r\n3. Scroll down to '....'\r\n4. See error\r\n\r\n<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->\r\n\r\n\r\n#### Code sample\r\n<!-- Ideally attach a minimal code sample to reproduce the decried issue.\r\nMinimal means having the shortest code but still preserving the bug. -->\r\n\r\n### Expected behavior\r\n\r\n<!-- A clear and concise description of what you expected to happen. -->\r\n\r\n### Environment\r\n\r\n - PyTorch Version (e.g., 1.0):\r\n - OS (e.g., Linux):\r\n - How you installed PyTorch (`conda`, `pip`, source):\r\n - Build command you used (if compiling from source):\r\n - Python version:\r\n - CUDA/cuDNN version:\r\n - GPU models and configuration:\r\n - Any other relevant information:\r\n\r\n### Additional context\r\n\r\n<!-- Add any other context about the problem here. -->\r\n\n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom flash.text.seq2seq.core.data import Seq2SeqData, Seq2SeqPostprocess, Seq2SeqPreprocess\n\n\nclass SummarizationData(Seq2SeqData):\n\n    preprocess_cls = Seq2SeqPreprocess\n    postprocess_cls = Seq2SeqPostprocess\n", "path": "flash/text/seq2seq/summarization/data.py"}, {"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Dict, List, Tuple\n\nimport numpy as np\nfrom torch import tensor\nfrom torchmetrics import Metric\n\nfrom flash.core.utilities.imports import _TEXT_AVAILABLE\nfrom flash.text.seq2seq.summarization.utils import add_newline_to_end_of_each_sentence\n\nif _TEXT_AVAILABLE:\n    from rouge_score import rouge_scorer, scoring\n    from rouge_score.scoring import AggregateScore, BootstrapAggregator, Score\nelse:\n    AggregateScore, Score, BootstrapAggregator = None, None, object\n\n\nclass RougeMetric(Metric):\n    \"\"\"\n    Metric used for automatic summarization. https://www.aclweb.org/anthology/W04-1013/\n\n    Example:\n\n        >>> target = \"Is your name John\".split()\n        >>> preds = \"My name is John\".split()\n        >>> rouge = RougeMetric()   # doctest: +SKIP\n        >>> from pprint import pprint\n        >>> pprint(rouge(preds, target))  # doctest: +NORMALIZE_WHITESPACE +SKIP\n        {'rouge1_fmeasure': 0.25,\n         'rouge1_precision': 0.25,\n         'rouge1_recall': 0.25,\n         'rouge2_fmeasure': 0.0,\n         'rouge2_precision': 0.0,\n         'rouge2_recall': 0.0,\n         'rougeL_fmeasure': 0.25,\n         'rougeL_precision': 0.25,\n         'rougeL_recall': 0.25,\n         'rougeLsum_fmeasure': 0.25,\n         'rougeLsum_precision': 0.25,\n         'rougeLsum_recall': 0.25}\n    \"\"\"\n\n    def __init__(\n        self,\n        rouge_newline_sep: bool = False,\n        use_stemmer: bool = False,\n        rouge_keys: Tuple[str] = (\"rouge1\", \"rouge2\", \"rougeL\", \"rougeLsum\"),\n    ):\n        super().__init__()\n        if not _TEXT_AVAILABLE:\n            raise ModuleNotFoundError(\"Please, pip install 'lightning-flash[text]'\")\n\n        self.rouge_newline_sep = rouge_newline_sep\n        self.rouge_keys = rouge_keys\n        self.use_stemmer = use_stemmer\n        self.aggregator = RougeBatchAggregator()\n        self.scorer = rouge_scorer.RougeScorer(rouge_keys, use_stemmer=self.use_stemmer)\n\n        for key in rouge_keys:\n            self.add_state(key, [])\n\n    def update(self, pred_lns: List[str], tgt_lns: List[str]):\n        for pred, tgt in zip(pred_lns, tgt_lns):\n            # rougeLsum expects \"\\n\" separated sentences within a summary\n            if self.rouge_newline_sep:\n                pred = add_newline_to_end_of_each_sentence(pred)\n                tgt = add_newline_to_end_of_each_sentence(tgt)\n            results = self.scorer.score(pred, tgt)\n            for key, score in results.items():\n                score = tensor([score.precision, score.recall, score.fmeasure])\n                getattr(self, key).append(score)\n\n    def compute(self) -> Dict[str, float]:\n        scores = {key: getattr(self, key) for key in self.rouge_keys}\n        self.aggregator.add_scores(scores)\n        result = self.aggregator.aggregate()\n        return format_rouge_results(result)\n\n    def __hash__(self):\n        # override to hash list objects.\n        # this is a bug in the upstream pytorch release.\n        hash_vals = [self.__class__.__name__]\n\n        for key in self._defaults.keys():\n            value = getattr(self, key)\n            if isinstance(value, list):\n                value = tuple(value)\n            hash_vals.append(value)\n\n        return hash(tuple(hash_vals))\n\n\nclass RougeBatchAggregator(BootstrapAggregator):\n    \"\"\"\n    Aggregates rouge scores and provides confidence intervals.\n    \"\"\"\n\n    def aggregate(self):\n        \"\"\"\n        Override function to wrap the final results in `Score` objects.\n        This is due to the scores being replaced with a list of torch tensors.\n        \"\"\"\n        result = {}\n        for score_type, scores in self._scores.items():\n            # Stack scores into a 2-d matrix of (sample, measure).\n            score_matrix = np.vstack(tuple(scores))\n            # Percentiles are returned as (interval, measure).\n            percentiles = self._bootstrap_resample(score_matrix)\n            # Extract the three intervals (low, mid, high).\n            intervals = tuple((Score(*percentiles[j, :]) for j in range(3)))\n            result[score_type] = AggregateScore(low=intervals[0], mid=intervals[1], high=intervals[2])\n        return result\n\n    def add_scores(self, scores):\n        self._scores = scores\n\n\ndef format_rouge_results(result: Dict[str, AggregateScore], decimal_places: int = 4) -> Dict[str, float]:\n    flattened_result = {}\n    for rouge_key, rouge_aggregate_score in result.items():\n        for stat in [\"precision\", \"recall\", \"fmeasure\"]:\n            mid = rouge_aggregate_score.mid\n            score = round(getattr(mid, stat), decimal_places)\n            flattened_result[f\"{rouge_key}_{stat}\"] = score\n    return flattened_result\n", "path": "flash/text/seq2seq/summarization/metric.py"}, {"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union\n\nimport pytorch_lightning as pl\nimport torch\n\nfrom flash.text.seq2seq.core.model import Seq2SeqTask\nfrom flash.text.seq2seq.summarization.metric import RougeMetric\n\n\nclass SummarizationTask(Seq2SeqTask):\n    \"\"\"Task for Seq2Seq Summarization.\n\n    Args:\n        backbone: backbone model to use for the task.\n        loss_fn: Loss function for training.\n        optimizer: Optimizer to use for training, defaults to `torch.optim.Adam`.\n        metrics: Metrics to compute for training and evaluation.\n        learning_rate: Learning rate to use for training, defaults to `3e-4`\n        val_target_max_length: Maximum length of targets in validation. Defaults to `128`\n        num_beams: Number of beams to use in validation when generating predictions. Defaults to `4`\n        use_stemmer: Whether Porter stemmer should be used to strip word suffixes to improve matching.\n        rouge_newline_sep: Add a new line at the beginning of each sentence in Rouge Metric calculation.\n    \"\"\"\n\n    def __init__(\n        self,\n        backbone: str = \"sshleifer/tiny-mbart\",\n        loss_fn: Optional[Union[Callable, Mapping, Sequence]] = None,\n        optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,\n        metrics: Union[pl.metrics.Metric, Mapping, Sequence, None] = None,\n        learning_rate: float = 5e-5,\n        val_target_max_length: Optional[int] = None,\n        num_beams: Optional[int] = 4,\n        use_stemmer: bool = True,\n        rouge_newline_sep: bool = True\n    ):\n        self.save_hyperparameters()\n        super().__init__(\n            backbone=backbone,\n            loss_fn=loss_fn,\n            optimizer=optimizer,\n            metrics=metrics,\n            learning_rate=learning_rate,\n            val_target_max_length=val_target_max_length,\n            num_beams=num_beams\n        )\n        self.rouge = RougeMetric(\n            rouge_newline_sep=rouge_newline_sep,\n            use_stemmer=use_stemmer,\n        )\n\n    @property\n    def task(self) -> str:\n        return \"summarization\"\n\n    def compute_metrics(self, generated_tokens: torch.Tensor, batch: Dict, prefix: str) -> None:\n        tgt_lns = self.tokenize_labels(batch[\"labels\"])\n        result = self.rouge(self._postprocess.uncollate(generated_tokens), tgt_lns)\n        self.log_dict(result, on_step=False, on_epoch=True)\n\n    def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):\n        \"\"\"\n        This function is used only for debugging usage with CI\n        \"\"\"\n        assert history[-1][\"val_f1\"] > 0.45\n", "path": "flash/text/seq2seq/summarization/model.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Callable, Dict, Optional, Union\n\nfrom flash.text.seq2seq.core.data import Seq2SeqData, Seq2SeqPostprocess, Seq2SeqPreprocess\n\n\nclass SummarizationPreprocess(Seq2SeqPreprocess):\n\n    def __init__(\n        self,\n        train_transform: Optional[Dict[str, Callable]] = None,\n        val_transform: Optional[Dict[str, Callable]] = None,\n        test_transform: Optional[Dict[str, Callable]] = None,\n        predict_transform: Optional[Dict[str, Callable]] = None,\n        backbone: str = \"sshleifer/distilbart-xsum-1-1\",\n        max_source_length: int = 128,\n        max_target_length: int = 128,\n        padding: Union[str, bool] = 'max_length'\n    ):\n        super().__init__(\n            train_transform=train_transform,\n            val_transform=val_transform,\n            test_transform=test_transform,\n            predict_transform=predict_transform,\n            backbone=backbone,\n            max_source_length=max_source_length,\n            max_target_length=max_target_length,\n            padding=padding,\n        )\n\n\nclass SummarizationData(Seq2SeqData):\n\n    preprocess_cls = SummarizationPreprocess\n    postprocess_cls = Seq2SeqPostprocess\n", "path": "flash/text/seq2seq/summarization/data.py"}, {"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Dict, List, Tuple\n\nimport numpy as np\nfrom torch import tensor\nfrom torchmetrics import Metric\n\nfrom flash.core.utilities.imports import _TEXT_AVAILABLE\nfrom flash.text.seq2seq.summarization.utils import add_newline_to_end_of_each_sentence\n\nif _TEXT_AVAILABLE:\n    from rouge_score import rouge_scorer\n    from rouge_score.scoring import AggregateScore, BootstrapAggregator, Score\nelse:\n    AggregateScore, Score, BootstrapAggregator = None, None, object\n\n\nclass RougeMetric(Metric):\n    \"\"\"\n    Metric used for automatic summarization. https://www.aclweb.org/anthology/W04-1013/\n\n    Example:\n\n        >>> target = \"Is your name John\".split()\n        >>> preds = \"My name is John\".split()\n        >>> rouge = RougeMetric()   # doctest: +SKIP\n        >>> from pprint import pprint\n        >>> pprint(rouge(preds, target))  # doctest: +NORMALIZE_WHITESPACE +SKIP\n        {'rouge1_fmeasure': 0.25,\n         'rouge1_precision': 0.25,\n         'rouge1_recall': 0.25,\n         'rouge2_fmeasure': 0.0,\n         'rouge2_precision': 0.0,\n         'rouge2_recall': 0.0,\n         'rougeL_fmeasure': 0.25,\n         'rougeL_precision': 0.25,\n         'rougeL_recall': 0.25,\n         'rougeLsum_fmeasure': 0.25,\n         'rougeLsum_precision': 0.25,\n         'rougeLsum_recall': 0.25}\n    \"\"\"\n\n    def __init__(\n        self,\n        rouge_newline_sep: bool = False,\n        use_stemmer: bool = False,\n        rouge_keys: Tuple[str] = (\"rouge1\", \"rouge2\", \"rougeL\", \"rougeLsum\"),\n    ):\n        super().__init__()\n        if not _TEXT_AVAILABLE:\n            raise ModuleNotFoundError(\"Please, pip install 'lightning-flash[text]'\")\n\n        self.rouge_newline_sep = rouge_newline_sep\n        self.rouge_keys = rouge_keys\n        self.use_stemmer = use_stemmer\n        self.aggregator = RougeBatchAggregator()\n        self.scorer = rouge_scorer.RougeScorer(rouge_keys, use_stemmer=self.use_stemmer)\n\n        for key in rouge_keys:\n            self.add_state(key, [])\n\n    def update(self, pred_lns: List[str], tgt_lns: List[str]):\n        for pred, tgt in zip(pred_lns, tgt_lns):\n            # rougeLsum expects \"\\n\" separated sentences within a summary\n            if self.rouge_newline_sep:\n                pred = add_newline_to_end_of_each_sentence(pred)\n                tgt = add_newline_to_end_of_each_sentence(tgt)\n            results = self.scorer.score(pred, tgt)\n            for key, score in results.items():\n                score = tensor([score.precision, score.recall, score.fmeasure])\n                getattr(self, key).append(score)\n\n    def compute(self) -> Dict[str, float]:\n        scores = {key: getattr(self, key) for key in self.rouge_keys}\n        self.aggregator.add_scores(scores)\n        result = self.aggregator.aggregate()\n        return format_rouge_results(result)\n\n    def __hash__(self):\n        # override to hash list objects.\n        # this is a bug in the upstream pytorch release.\n        hash_vals = [self.__class__.__name__]\n\n        for key in self._defaults.keys():\n            value = getattr(self, key)\n            if isinstance(value, list):\n                value = tuple(value)\n            hash_vals.append(value)\n\n        return hash(tuple(hash_vals))\n\n\nclass RougeBatchAggregator(BootstrapAggregator):\n    \"\"\"\n    Aggregates rouge scores and provides confidence intervals.\n    \"\"\"\n\n    def aggregate(self):\n        \"\"\"\n        Override function to wrap the final results in `Score` objects.\n        This is due to the scores being replaced with a list of torch tensors.\n        \"\"\"\n        result = {}\n        for score_type, scores in self._scores.items():\n            # Stack scores into a 2-d matrix of (sample, measure).\n            score_matrix = np.vstack(tuple(scores))\n            # Percentiles are returned as (interval, measure).\n            percentiles = self._bootstrap_resample(score_matrix)\n            # Extract the three intervals (low, mid, high).\n            intervals = tuple((Score(*percentiles[j, :]) for j in range(3)))\n            result[score_type] = AggregateScore(low=intervals[0], mid=intervals[1], high=intervals[2])\n        return result\n\n    def add_scores(self, scores):\n        self._scores = scores\n\n\ndef format_rouge_results(result: Dict[str, AggregateScore], decimal_places: int = 4) -> Dict[str, float]:\n    flattened_result = {}\n    for rouge_key, rouge_aggregate_score in result.items():\n        for stat in [\"precision\", \"recall\", \"fmeasure\"]:\n            mid = rouge_aggregate_score.mid\n            score = round(getattr(mid, stat), decimal_places)\n            flattened_result[f\"{rouge_key}_{stat}\"] = score\n    return flattened_result\n", "path": "flash/text/seq2seq/summarization/metric.py"}, {"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union\n\nimport pytorch_lightning as pl\nimport torch\n\nfrom flash.text.seq2seq.core.model import Seq2SeqTask\nfrom flash.text.seq2seq.summarization.metric import RougeMetric\n\n\nclass SummarizationTask(Seq2SeqTask):\n    \"\"\"Task for Seq2Seq Summarization.\n\n    Args:\n        backbone: backbone model to use for the task.\n        loss_fn: Loss function for training.\n        optimizer: Optimizer to use for training, defaults to `torch.optim.Adam`.\n        metrics: Metrics to compute for training and evaluation.\n        learning_rate: Learning rate to use for training, defaults to `3e-4`\n        val_target_max_length: Maximum length of targets in validation. Defaults to `128`\n        num_beams: Number of beams to use in validation when generating predictions. Defaults to `4`\n        use_stemmer: Whether Porter stemmer should be used to strip word suffixes to improve matching.\n        rouge_newline_sep: Add a new line at the beginning of each sentence in Rouge Metric calculation.\n    \"\"\"\n\n    def __init__(\n        self,\n        backbone: str = \"sshleifer/distilbart-xsum-1-1\",\n        loss_fn: Optional[Union[Callable, Mapping, Sequence]] = None,\n        optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,\n        metrics: Union[pl.metrics.Metric, Mapping, Sequence, None] = None,\n        learning_rate: float = 1e-5,\n        val_target_max_length: Optional[int] = None,\n        num_beams: Optional[int] = 4,\n        use_stemmer: bool = True,\n        rouge_newline_sep: bool = True\n    ):\n        self.save_hyperparameters()\n        super().__init__(\n            backbone=backbone,\n            loss_fn=loss_fn,\n            optimizer=optimizer,\n            metrics=metrics,\n            learning_rate=learning_rate,\n            val_target_max_length=val_target_max_length,\n            num_beams=num_beams\n        )\n        self.rouge = RougeMetric(\n            rouge_newline_sep=rouge_newline_sep,\n            use_stemmer=use_stemmer,\n        )\n\n    @property\n    def task(self) -> str:\n        return \"summarization\"\n\n    def compute_metrics(self, generated_tokens: torch.Tensor, batch: Dict, prefix: str) -> None:\n        tgt_lns = self.tokenize_labels(batch[\"labels\"])\n        result = self.rouge(self._postprocess.uncollate(generated_tokens), tgt_lns)\n        self.log_dict(result, on_step=False, on_epoch=True, prog_bar=True)\n\n    def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):\n        \"\"\"\n        This function is used only for debugging usage with CI\n        \"\"\"\n        assert history[-1][\"rouge1_recall\"] > 0.2\n", "path": "flash/text/seq2seq/summarization/model.py"}]}