def prepare_for_training(
        self,
        framework: Union[Framework, str] = "transformers",
        settings: Union[TextClassificationSettings, TokenClassificationSettings] = None,
        lang: Optional["spacy.Language"] = None,
        train_size: Optional[float] = 1,
        test_size: Optional[float] = None,
        seed: Optional[int] = None,
    ) -> Union[
        "datasets.Dataset",
        "spacy.tokens.DocBin",
        Tuple["spacy.tokens.DocBin", "spacy.tokens.DocBin"],
        Tuple["pandas.DataFrame", "pandas.DataFrame"],
    ]:
        """Prepares the dataset for training.

        This will return a ``datasets.Dataset`` with all columns returned by ``to_datasets`` method
        and an additional  *ner_tags* column:
            - Records without an annotation are removed.
            - The *ner_tags* column corresponds to the iob tags sequences for annotations of the records
            - The iob tags are transformed to integers.

        Args:
            framework: A string|enum specifying the framework for the training.
                "transformers" and "spacy" are currently supported. Default: `transformers`
            lang: The spacy nlp Language pipeline used to process the dataset. (Only for spacy framework)
            train_size: The size of the training set. If float, should be between 0.0 and 1.0 and represent the
                proportion of the dataset to include in the train split.
            test_size: The size of the test set. If float, should be between 0.0 and 1.0 and represent the
                proportion of the dataset to include in the test split.
            seed: Random state.

        Returns:
            A datasets Dataset with a *ner_tags* or a *label* column and several *inputs* columns.
            returned by ``to_datasets`` for "transformers" framework or a spaCy DocBin for "spacy" framework.

        Examples:
            >>> import argilla_v1 as rg
            >>> rb_dataset = rg.DatasetForTokenClassification([
            ...     rg.TokenClassificationRecord(
            ...         text="The text",
            ...         tokens=["The", "text"],
            ...         annotation=[("TAG", 0, 2)],
            ...     )
            ... ])
            >>> rb_dataset.prepare_for_training().features
            {'text': Value(dtype='string'),
             'tokens': Sequence(feature=Value(dtype='string'), length=-1),
             'prediction': Value(dtype='null'),
             'prediction_agent': Value(dtype='null'),
             'annotation': [{'end': Value(dtype='int64'),
               'label': Value(dtype='string'),
               'start': Value(dtype='int64')}],
             'annotation_agent': Value(dtype='null'),
             'id': Value(dtype='null'),
             'metadata': Value(dtype='null'),
             'status': Value(dtype='string'),
             'event_timestamp': Value(dtype='null'),
             'metrics': Value(dtype='null'),
             'ner_tags': [ClassLabel(num_classes=3, names=['O', 'B-TAG', 'I-TAG'])]}

            >>> import argilla_v1 as rg
            >>> rb_dataset = rg.DatasetForTextClassification([
            ...     rg.TextClassificationRecord(
            ...         inputs={"header": "my header", "content": "my content"},
            ...         annotation="SPAM",
            ...     )
            ... ])
            >>> rb_dataset.prepare_for_training().features
            {'header': Value(dtype='string'),
             'content': Value(dtype='string'),
             'label': ClassLabel(num_classes=1, names=['SPAM'])}

        """
        if self._RECORD_TYPE == TextClassificationRecord:
            if settings is None:
                self._SETTINGS = self._infer_settings_from_records()
            elif isinstance(settings, TextClassificationSettings):
                self._SETTINGS = settings
            else:
                raise ValueError("settings must be TextClassificationSettings for TextClassificationRecord")
        elif self._RECORD_TYPE == TokenClassificationRecord:
            if settings is None:
                self._SETTINGS = self._infer_settings_from_records()
            elif isinstance(settings, TokenClassificationSettings):
                self._SETTINGS = settings
            else:
                raise ValueError("settings must be TokenClassificationSettings for TokenClassificationRecord")
        else:
            self._SETTINGS = settings

        if train_size is None:
            train_size = 1
        if test_size is None:
            test_size = 1 - train_size

        # check if all numbers are larger than 0
        assert [abs(train_size), abs(test_size)] == [train_size, test_size], ValueError(
            "`train_size` and `test_size` must be larger than 0."
        )

        # check if train sizes sum up to 1
        assert (train_size + test_size) == 1, ValueError("`train_size` and `test_size` must sum to 1.")

        # check for annotations
        assert any([rec.annotation for rec in self._records]), ValueError("Dataset has no annotations.")

        if test_size == 0:
            test_size = None

        # shuffle records
        shuffled_records = self._records.copy()
        seed = seed or random.randint(42, 1984)
        random.Random(seed).shuffle(shuffled_records)

        # turn the string into a Framework instance and trigger error if str is not valid
        if isinstance(framework, str):
            framework = Framework(framework)

        # prepare for training for the right method
        if framework in [
            Framework.TRANSFORMERS,
            Framework.SETFIT,
            Framework.SPAN_MARKER,
            Framework.PEFT,
        ]:
            return self._prepare_for_training_with_transformers(train_size=train_size, test_size=test_size, seed=seed)
        elif framework in [Framework.SPACY, Framework.SPACY_TRANSFORMERS] and lang is None:
            raise ValueError(
                "Please provide a `spaCy` language model to prepare the dataset for"
                " training with the `spaCy`/`spaCy-transformers` framework."
            )
        elif framework in [Framework.SPACY, Framework.SPACY_TRANSFORMERS, Framework.SPARK_NLP, Framework.OPENAI]:
            if train_size and test_size:
                require_dependencies("scikit-learn")
                from sklearn.model_selection import train_test_split

                records_train, records_test = train_test_split(
                    shuffled_records,
                    train_size=train_size,
                    shuffle=False,
                    random_state=seed,
                )
                if framework in [Framework.SPACY, Framework.SPACY_TRANSFORMERS]:
                    train_docbin = self._prepare_for_training_with_spacy(nlp=lang, records=records_train)
                    test_docbin = self._prepare_for_training_with_spacy(nlp=lang, records=records_test)
                    return train_docbin, test_docbin
                elif framework is Framework.SPARK_NLP:
                    train_df = self._prepare_for_training_with_spark_nlp(records_train)
                    test_df = self._prepare_for_training_with_spark_nlp(records_test)
                    return train_df, test_df
                else:
                    train_jsonl = self._prepare_for_training_with_openai(records=records_train)
                    test_jsonl = self._prepare_for_training_with_openai(records=records_test)
                    return train_jsonl, test_jsonl
            else:
                if framework in [Framework.SPACY, Framework.SPACY_TRANSFORMERS]:
                    return self._prepare_for_training_with_spacy(nlp=lang, records=shuffled_records)
                elif framework is Framework.SPARK_NLP:
                    return self._prepare_for_training_with_spark_nlp(records=shuffled_records)
                elif framework is Framework.OPENAI:
                    return self._prepare_for_training_with_openai(records=shuffled_records)
                else:
                    raise NotImplementedError(
                        f"Framework {framework} is not supported. Choose from: {[e.value for e in Framework]}"
                    )
        else:
            raise NotImplementedError(
                f"Framework {framework} is not supported. Choose from: {[e.value for e in Framework]}"
            )
Prepares the dataset for training.

        This will return a ``datasets.Dataset`` with all columns returned by ``to_datasets`` method
        and an additional  *ner_tags* column:
            - Records without an annotation are removed.
            - The *ner_tags* column corresponds to the iob tags sequences for annotations of the records
            - The iob tags are transformed to integers.

        Args:
            framework: A string|enum specifying the framework for the training.
                "transformers" and "spacy" are currently supported. Default: `transformers`
            lang: The spacy nlp Language pipeline used to process the dataset. (Only for spacy framework)
            train_size: The size of the training set. If float, should be between 0.0 and 1.0 and represent the
                proportion of the dataset to include in the train split.
            test_size: The size of the test set. If float, should be between 0.0 and 1.0 and represent the
                proportion of the dataset to include in the test split.
            seed: Random state.

        Returns:
            A datasets Dataset with a *ner_tags* or a *label* column and several *inputs* columns.
            returned by ``to_datasets`` for "transformers" framework or a spaCy DocBin for "spacy" framework.

        Examples:
            >>> import argilla_v1 as rg
            >>> rb_dataset = rg.DatasetForTokenClassification([
            ...     rg.TokenClassificationRecord(
            ...         text="The text",
            ...         tokens=["The", "text"],
            ...         annotation=[("TAG", 0, 2)],
            ...     )
            ... ])
            >>> rb_dataset.prepare_for_training().features
            {'text': Value(dtype='string'),
             'tokens': Sequence(feature=Value(dtype='string'), length=-1),
             'prediction': Value(dtype='null'),
             'prediction_agent': Value(dtype='null'),
             'annotation': [{'end': Value(dtype='int64'),
               'label': Value(dtype='string'),
               'start': Value(dtype='int64')}],
             'annotation_agent': Value(dtype='null'),
             'id': Value(dtype='null'),
             'metadata': Value(dtype='null'),
             'status': Value(dtype='string'),
             'event_timestamp': Value(dtype='null'),
             'metrics': Value(dtype='null'),
             'ner_tags': [ClassLabel(num_classes=3, names=['O', 'B-TAG', 'I-TAG'])]}

            >>> import argilla_v1 as rg
            >>> rb_dataset = rg.DatasetForTextClassification([
            ...     rg.TextClassificationRecord(
            ...         inputs={"header": "my header", "content": "my content"},
            ...         annotation="SPAM",
            ...     )
            ... ])
            >>> rb_dataset.prepare_for_training().features
            {'header': Value(dtype='string'),
             'content': Value(dtype='string'),
             'label': ClassLabel(num_classes=1, names=['SPAM'])}

        
argilla-v1/src/argilla_v1/client/datasets.py