[wip] adding SciTLDR dataset

catwalk/models/t5.py

@@ -70,6 +70,29 @@ def _predict_prompt(self, task: Task, instances: Sequence[Dict[str, Any]], batch
                             "rouge": ([prediction], [[target]])
                         }
 
+    def _predict_summarization(self, task: Task, instances: Sequence[Dict[str, Any]], batch_size: int = 32) -> Iterator[Dict[str, Any]]:
+        examples = MappedSequence(task.instance_conversions[InstanceFormat.HF_SUMMARIZATION], instances)
+
+        model = self.get_model().eval()
+        tokenizer = self.get_tokenizer()
+        tokenizer.model_max_length = model.config.n_positions
+
+        with torch.inference_mode():
+            with Tqdm.tqdm(examples, desc="Processing instances") as examples_tqdm:
+                for batch in more_itertools.chunked(examples_tqdm, batch_size):
+                    model_input = tokenizer([f"summarize:{i.source}" for i in batch],
+                                            truncation="only_first",
+                                            padding="longest",
+                                            return_tensors="pt")
+
+                    model_output = model.generate(**model_input, max_new_tokens=50)
+                    model_output = tokenizer.batch_decode(model_output, clean_up_tokenization_spaces=True, skip_special_tokens=True)
+                    for instance, prediction in zip(batch, model_output):
+                        target = instance.target
+                        yield {
+                            "rouge": ([prediction], [target])
+                        }
+
     def predict(  # type: ignore
         self,
         task: Task,
@@ -81,6 +104,8 @@ def predict(  # type: ignore
             return self._predict_prompt(task, instances, batch_size=batch_size)
         elif task.has_instance_conversion(InstanceFormat.HF_QA):
             return self._predict_qa(task, instances, batch_size=batch_size)
+        elif task.has_instance_conversion(InstanceFormat.HF_SUMMARIZATION):
+            return self._predict_summarization(task, instances, batch_size=batch_size)
 
         raise UnsupportedTaskError(self, task)
 

catwalk/task.py

@@ -24,6 +24,9 @@
     "squad_metrics": torchmetrics.SQuAD,
 }
 
+SUMMARIZATION_METRICS = {
+    'summarization_metrics': torchmetrics.text.ROUGEScore
+}
 
 try:
     from functools import cache as memoize
@@ -59,6 +62,7 @@ class InstanceFormat(Enum):
     HF_MC = 2
     HF_QA = 8
     HF_CLASSIFICATION = 10
+    HF_SUMMARIZATION = 11
     ELEUTHER_DOC = 3
     ELEUTHER_CONTEXT = 4
     ELEUTHER_REQUESTS = 5

catwalk/tasks/__init__.py

@@ -2,11 +2,12 @@
 
 import datasets
 
-from catwalk.task import InstanceFormat, ENTAILMENT_METRICS, QA_METRICS, Task, \
+from catwalk.task import InstanceFormat, ENTAILMENT_METRICS, QA_METRICS, SUMMARIZATION_METRICS, Task, \
     classification_metrics, BINARY_CLASSIFICATION_METRICS, mc_metrics, PERPLEXITY_METRICS
 from catwalk.tasks.eleuther import EleutherTask, RaceEleutherTask, EleutherTaskWithRenamedSplits, \
     EleutherClassificationTask, EleutherClassificationTaskWithRenamedSplits
-from catwalk.tasks.huggingface import hfmc_conversion, HFDatasetsTask, hfqa_conversion, hfclassification_conversion
+from catwalk.tasks.huggingface import hfmc_conversion, HFDatasetsTask, hfqa_conversion, \
+    hfclassification_conversion, hfsummarization_conversion
 from catwalk.tasks.p3 import P3Task
 from catwalk.tasks.raft import RaftTask
 from catwalk.tasks.metaicl import MetaICLTask
@@ -449,6 +450,12 @@
     "metaicl::numer_sense": MetaICLTask("numer_sense").add_metrics(classification_metrics(12)),
     "metaicl::race-high": MetaICLTask("race-high").add_metrics(mc_metrics(4)),
     "metaicl::commonsense_qa": MetaICLTask("commonsense_qa").add_metrics(mc_metrics(5)),
+
+    # Summarization
+    "scitldr": HFDatasetsTask("scitldr").add_instance_conversion(
+        InstanceFormat.HF_SUMMARIZATION,
+        hfsummarization_conversion()
+    ).add_metrics(SUMMARIZATION_METRICS),
 }
 
 for config in datasets.get_dataset_config_names("bigscience/P3"):
@@ -575,6 +582,9 @@
         "metaicl::tweet_eval-hate",
         "metaicl::tweet_eval-stance_atheism",
         "metaicl::tweet_eval-stance_feminist"
+    },
+    "s2": {
+        "scitldr"
     }
 }
 

catwalk/tasks/huggingface.py

@@ -222,3 +222,31 @@ def hfclassification_conversion(
 ) -> InstanceConversion:
     # We're doing this in this stupid way because this makes the conversion function picklable.
     return functools.partial(hfclassification_convert, **kwargs)
+
+
+
+@dataclass
+class HFSummarizationInstance:
+    id: Optional[str]
+    source: str
+    target: str
+
+
+def hfsummarization_convert(
+    instance: Dict[str, Any],
+    *,
+    source_field: str = "source",
+    target_field: str = "target",
+    id_field: Optional[str] = None,
+) -> HFSummarizationInstance:
+    return HFSummarizationInstance(
+        id=str(get_from_dict(instance, id_field)) if id_field else None,
+        source=str(get_from_dict(instance, source_field)),
+        target=str(get_from_dict(instance, target_field)))
+
+
+def hfsummarization_conversion(
+    **kwargs,
+) -> InstanceConversion:
+    # We're doing this in this stupid way because this makes the conversion function picklable.
+    return functools.partial(hfsummarization_convert, **kwargs)