add task categories

src/synthetic_dataset_generator/_distiset.py

@@ -81,6 +81,15 @@ class CustomDistisetWithAdditionalTag(distilabel.distiset.Distiset):
                 dataset[0] if not isinstance(dataset, dict) else dataset["train"][0]
             )
 
+        keys = list(sample_records.keys())
+        if len(keys) != 2 or not (
+            ("label" in keys and "text" in keys)
+            or ("labels" in keys and "text" in keys)
+        ):
+            task_categories = ["text-classification"]
+        elif "prompt" in keys or "messages" in keys:
+            task_categories = ["text-generation", "text2text-generation"]
+
         readme_metadata = {}
         if repo_id and token:
             readme_metadata = self._extract_readme_metadata(repo_id, token)
@@ -90,6 +99,7 @@ class CustomDistisetWithAdditionalTag(distilabel.distiset.Distiset):
             "size_categories": size_categories_parser(
                 max(len(dataset) for dataset in self.values())
             ),
+            "task_categories": task_categories,
             "tags": [
                 "synthetic",
                 "distilabel",