Feature/Added support for CSV augmentations in text classification task

langtest/datahandler/datasource.py

@@ -398,15 +398,22 @@ def export_data(self, data: List[Sample], output_path: str):
         """
         temp_id = None
         otext = ""
-        for i in data:
-            if isinstance(i, NEROutput):
+        if self.task == "ner":
+            for i in data:
                 text, temp_id = Formatter.process(i, output_format="csv", temp_id=temp_id)
-            else:
-                text = Formatter.process(i, output_format="csv")
-            otext += text
+                otext += text
 
-        with open(output_path, "wb") as fwriter:
-            fwriter.write(bytes(otext, encoding="utf-8"))
+            with open(output_path, "wb") as fwriter:
+                fwriter.write(bytes(otext, encoding="utf-8"))
+
+        elif self.task == "text-classification":
+            rows = []
+            for s in data:
+                row = Formatter.process(s, output_format="csv")
+                rows.append(row)
+
+            df = pd.DataFrame(rows, columns=list(self.COLUMN_NAMES.keys()))
+            df.to_csv(output_path, index=False, encoding="utf-8")
 
     @staticmethod
     def _find_delimiter(file_path: str) -> property:
@@ -825,12 +832,15 @@ def export_data(self, data: List[Sample], output_path: str):
             output_path (str):
                 Path to save the data to.
         """
-        with open(output_path, "w") as file:
-            csv_writer = csv.writer(file)
-            csv_writer.writerow(list(self.COLUMN_NAMES["text-classification"].keys()))
-            for s in data:
-                row = self._sample_to_row(s)
-                csv_writer.writerow(row)
+        rows = []
+        for s in data:
+            row = Formatter.process(s, output_format="csv")
+            rows.append(row)
+
+        df = pd.DataFrame(
+            rows, columns=list(self.COLUMN_NAMES["text-classification"].keys())
+        )
+        df.to_csv(output_path, index=False, encoding="utf-8")
 
     def _row_to_sample_classification(self, data_row: Dict[str, str]) -> Sample:
         """
@@ -868,19 +878,3 @@ def _row_to_sample_classification(self, data_row: Dict[str, str]) -> Sample:
             original=original,
             expected_results=SequenceClassificationOutput(predictions=[label]),
         )
-
-    @staticmethod
-    def _sample_to_row(s: Sample) -> List[str]:
-        """
-        Convert a Sample object into a row for exporting.
-
-        Args:
-            s (Sample):
-                Sample object to convert.
-
-        Returns:
-            List[str]:
-                Row formatted as a list of strings.
-        """
-        row = [s.original, s.expected_results.predictions[0].label]
-        return row

langtest/datahandler/format.py

@@ -90,24 +90,23 @@ class SequenceClassificationOutputFormatter(BaseFormatter, ABC):
     """
 
     @staticmethod
-    def to_csv(sample: Sample, delimiter: str = ",") -> str:
-        """Converts a custom type to a CSV string.
+    def to_csv(sample: Sample) -> str:
+        """
+        Convert a Sample object into a row for exporting.
 
         Args:
-            sample (Sample):
-                The input sample containing the `SequenceClassificationOutput` object to convert.
-            delimiter (str):
-                The delimiter character to use in the CSV string.
+            Sample :
+                Sample object to convert.
 
         Returns:
-            str: The CSV string representation of the `SequenceClassificationOutput` object.
+            List[str]:
+                Row formatted as a list of strings.
         """
-        original = sample.original
-        test_case = sample.test_case
-        if test_case:
-            return f"{test_case}{delimiter}{sample.expected_results.to_str_list()[0]}\n"
+        if sample.test_case:
+            row = [sample.test_case, sample.expected_results.predictions[0].label]
         else:
-            return f"{original}{delimiter}{sample.expected_results.to_str_list()[0]}\n"
+            row = [sample.original, sample.expected_results.predictions[0].label]
+        return row
 
 
 class NEROutputFormatter(BaseFormatter):