output in dataset

tasks.py

@@ -62,6 +62,7 @@ class Task:
     metric_name: str | tuple[str, str] = ("sustech/tlem", "mmlu")
     input_column: str = "question"
     label_column: str = ""
+    output_column: str = "generated_text"
     prompt: Optional[Callable | str] = None
     few_shot: int = 0
     few_shot_from: Optional[str] = None
@@ -85,7 +86,6 @@ class Task:
                 )
             }
         self.label_column = self.label_column or self.input_column
-        self.outputs = []
 
     def __eq__(self, __value: object) -> bool:
         return self.name == __value.name
@@ -98,6 +98,10 @@ class Task:
     def labels(self):
         return self.dataset[self.label_column]
 
+    @cached_property
+    def outputs(self):
+        return self.dataset[self.output_column]
+
     @cached_property
     def dataset(self):
         ds = (
@@ -160,20 +164,29 @@ class Task:
         # logging.info(f"{self.name}:{results}")
         return results
 
-    # @cache
     def run(
         self,
         pipeline,
     ):
-        self.outputs = self.outputs or pipeline(self.samples)
+        if self.output_column not in self.dataset.column_names:
+            self.dataset = self.dataset.add_column(
+                self.output_column, pipeline(self.samples)
+            )
 
         return self.result
 
     async def arun(self, pipeline):
-        self.outputs = self.outputs or await pipeline(self.samples)
+        self.dataset = self.dataset.add_column(
+            self.output_column, await pipeline(self.samples)
+        )
 
         return self.result
 
+    def save(self, path):
+        self.dataset.select_columns(
+            [self.input_column, self.output_column, self.label_column]
+        ).save_to_disk(path)
+
 
 def multichoice(responses: Any, references: list[str]):
     if isinstance(responses[0], str):