def predict(
        self, text: Union[List[List[str]], Tuple[str, List[str]]], as_argilla_records: bool = False, **kwargs
    ) -> List[float]:
        """
        Predicts the similarity of the sentences.

        Args:
            text: The sentences to obtain the similarity from.
                Allowed inputs are:
                - A list with a single sentence (as a string) and a list of sentences to compare against.
                - A list with pair of sentences.
            as_argilla_records: If True, the prediction will be returned as an Argilla record. If
                False, the prediction will be returned as a string. Defaults to True

        Returns:
            A list of predicted similarities.
        """
        if isinstance(text[0], str):
            sentences = [(text[0], sentence) for sentence in text[1]]
        else:
            sentences = text

        if self._cross_encoder:
            prediction = list(self._trainer.predict(sentences, **kwargs))

        else:
            from sentence_transformers.util import cos_sim

            # Under this model we could potentially obtain the similarity of all vs all
            # the input sentences, but maybe is more intuitive to return only the similarities
            # between pairs as we do with the CrossEncoder predictions.
            sources, targets = zip(*sentences)
            embeds_source = self._trainer.encode(sources, **kwargs)
            embeds_target = self._trainer.encode(targets, **kwargs)
            prediction = list(cos_sim(embeds_source, embeds_target).numpy()[0])

        if as_argilla_records:
            raise NotImplementedError("To be done")

        return prediction

        Predicts the similarity of the sentences.

        Args:
            text: The sentences to obtain the similarity from.
                Allowed inputs are:
                - A list with a single sentence (as a string) and a list of sentences to compare against.
                - A list with pair of sentences.
            as_argilla_records: If True, the prediction will be returned as an Argilla record. If
                False, the prediction will be returned as a string. Defaults to True

        Returns:
            A list of predicted similarities.
        
argilla-v1/src/argilla_v1/client/feedback/training/frameworks/sentence_transformers.py