def predict(
        self, text: Union[List[str], str], as_argilla_records: bool = True, **kwargs
    ) -> Union[Dict[str, Any], List[Dict[str, Any]], BaseModel, List[BaseModel]]:
        """Predict the labels for the given text using the trained pipeline.

        Args:
            text: A `str` or a `List[str]` with the text to predict the labels for.
            as_argilla_records:
                A `bool` indicating whether to return the predictions as `argilla` records
                or as `dicts`. Defaults to True.

        Returns:
            Either a `dict`, `BaseModel` (if `as_argilla_records` is True) or a `List[dict]`,
            `List[BaseModel]` (if `as_argilla_records` is True) with the predictions.
        """
        if self.trainer_model is None:
            self._logger.warning("Using model without fine-tuning.")
            self.init_model()

        str_input = False
        if isinstance(text, str):
            text = [text]
            str_input = True

        formatted_prediction = []
        docs = self.trainer_model.pipe(text, **kwargs)
        if as_argilla_records:
            for doc in docs:
                if "ner" in self._spacy_pipeline_components:
                    entities = [(ent.label_, ent.start_char, ent.end_char) for ent in doc.ents]
                    pred = {
                        "text": doc.text,
                        "tokens": [t.text for t in doc],
                        "prediction": entities,
                    }
                    pred = self._record_class(**pred)
                elif any([p in self._spacy_pipeline_components for p in ["textcat", "textcat_multilabel"]]):
                    pred = {
                        "text": doc.text,
                        "prediction": [(k, v) for k, v in doc.cats.items()],
                    }
                    pred = self._record_class(**pred, multi_label=self._multi_label)
                else:
                    continue
                formatted_prediction.append(pred)
        else:
            formatted_prediction = docs

        if str_input:
            formatted_prediction = list(formatted_prediction)[0]
        return formatted_prediction
Predict the labels for the given text using the trained pipeline.

        Args:
            text: A `str` or a `List[str]` with the text to predict the labels for.
            as_argilla_records:
                A `bool` indicating whether to return the predictions as `argilla` records
                or as `dicts`. Defaults to True.

        Returns:
            Either a `dict`, `BaseModel` (if `as_argilla_records` is True) or a `List[dict]`,
            `List[BaseModel]` (if `as_argilla_records` is True) with the predictions.
        
argilla-v1/src/argilla_v1/training/spacy.py