Add summarization template

src/datasets/tasks/__init__.py

@@ -3,17 +3,19 @@
 from ..utils.logging import get_logger
 from .base import TaskTemplate
 from .question_answering import QuestionAnsweringExtractive
+from .summarization import Summarization
 from .text_classification import TextClassification
 
 
-__all__ = ["TaskTemplate", "QuestionAnsweringExtractive", "TextClassification"]
+__all__ = ["TaskTemplate", "QuestionAnsweringExtractive", "TextClassification", "Summarization"]
 
 logger = get_logger(__name__)
 
 
 NAME2TEMPLATE = {
     QuestionAnsweringExtractive.task: QuestionAnsweringExtractive,
     TextClassification.task: TextClassification,
+    Summarization.task: Summarization,
 }
 
 

src/datasets/tasks/summarization.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import ClassVar, Dict
+
+from ..features import Features, Value
+from .base import TaskTemplate
+
+
+@dataclass(frozen=True)
+class Summarization(TaskTemplate):
+    # `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
+    task: str = "summarization"
+    input_schema: ClassVar[Features] = Features({"text": Value("string")})
+    label_schema: ClassVar[Features] = Features({"summary": Value("string")})
+    text_column: str = "text"
+    summary_column: str = "summary"
+
+    @property
+    def column_mapping(self) -> Dict[str, str]:
+        return {self.text_column: "text", self.summary_column: "summary"}