Fix docstrings and add to HTML build

Author: lewtun

docs/source/package_reference/main_classes.rst

@@ -32,7 +32,7 @@ The base class :class:`datasets.Dataset` implements a Dataset backed by an Apach
         info, split, builder_name, citation, config_name, dataset_size,
         description, download_checksums, download_size, features, homepage,
         license, size_in_bytes, supervised_keys, version,
-        from_csv, from_json, from_text,
+        from_csv, from_json, from_text, prepare_for_task,
 
 .. autofunction:: datasets.concatenate_datasets
 
@@ -54,7 +54,7 @@ It also has dataset transform methods like map or filter, to process all the spl
         flatten_, cast_, remove_columns_, rename_column_,
         flatten, cast, remove_columns, rename_column, class_encode_column,
         save_to_disk, load_from_disk,
-        from_csv, from_json, from_text,
+        from_csv, from_json, from_text, prepare_for_task
 
 
 ``Features``

src/datasets/arrow_dataset.py

@@ -1391,12 +1391,12 @@ def prepare_for_task(self, task: Union[str, TaskTemplate]) -> "Dataset":
         Casts :attr:`datasets.DatasetInfo.features` according to a task-specific schema.
 
         Args:
-            task (:obj:`Union[str, TaskTemplate]`): The task to prepare the dataset for during training and evaluation. If `str`, supported tasks include:
+            task (:obj:`Union[str, TaskTemplate]`): The task to prepare the dataset for during training and evaluation. If :obj:`str`, supported tasks include:
 
-                - :obj:`"text-clasification"`
+                - :obj:`"text-classification"`
                 - :obj:`"question-answering"`
 
-            If `TaskTemplate`, must be one of the task templates in `datasets.tasks`.
+                If :obj:`TaskTemplate`, must be one of the task templates in :obj:`datasets.tasks`.
         """
         # TODO(lewtun): Add support for casting nested features like answers.text and answers.answer_start in SQuAD
         if isinstance(task, str):

src/datasets/dataset_dict.py

@@ -13,6 +13,7 @@
 from .features import Features
 from .filesystems import extract_path_from_uri, is_remote_filesystem
 from .table import Table
+from .tasks import TaskTemplate
 from .utils.deprecation_utils import deprecated
 from .utils.typing import PathLike
 
@@ -791,16 +792,18 @@ def from_text(
             path_or_paths, features=features, cache_dir=cache_dir, keep_in_memory=keep_in_memory, **kwargs
         ).read()
 
-    def prepare_for_task(self, task: str):
+    def prepare_for_task(self, task: Union[str, TaskTemplate]):
         """Prepare a dataset for the given task.
 
         Casts :attr:`datasets.DatasetInfo.features` according to a task-specific schema.
 
         Args:
-            task (``str``): The task to prepare the dataset for during training and evaluation. Supported tasks include:
+            task (:obj:`Union[str, TaskTemplate]`): The task to prepare the dataset for during training and evaluation. If :obj:`str`, supported tasks include:
 
-                - :obj:`"text-clasification"`
+                - :obj:`"text-classification"`
                 - :obj:`"question-answering"`
+
+                If :obj:`TaskTemplate`, must be one of the task templates in :obj:`datasets.tasks`.
         """
         self._check_values_type()
         return DatasetDict({k: dataset.prepare_for_task(task=task) for k, dataset in self.items()})