def extend_matrix(
        self,
        thresholds: Union[List[float], np.ndarray],
        embeddings: Optional[np.ndarray] = None,
        gpu: bool = False,
    ):
        """Extends the weak label matrix through embeddings according to the similarity thresholds for each rule.

        Implementation based on `Epoxy <https://github.com/HazyResearch/epoxy>`__.

        Args:
            thresholds: An array of thresholds between 0.0 and 1.0, one for each column of the weak labels matrix.
                Each one stands for the minimum cosine similarity between two sentences for a rule to be extended.
            embeddings: Embeddings for each row of the weak label matrix.
                If not provided, we will use the ones from the last ``WeakLabels.extend_matrix()`` call.
            gpu: If True, perform FAISS similarity queries on GPU.

        Examples:
            >>> # Choose any model to generate the embeddings.
            >>> from sentence_transformers import SentenceTransformer
            >>> model = SentenceTransformer('all-mpnet-base-v2', device='cuda')
            >>>
            >>> # Generate the embeddings and set the thresholds.
            >>> weak_labels = {class_name}(dataset="my_dataset")
            >>> embeddings = np.array([ model.encode(rec.text) for rec in weak_labels.records() ])
            >>> thresholds = [0.6] * len(weak_labels.rules)
            >>>
            >>> # Extend the weak labels matrix.
            >>> weak_labels.extend_matrix(thresholds, embeddings)
            >>>
            >>> # Calling the method below will now retrieve the extended matrix.
            >>> weak_labels.matrix()
            >>>
            >>> # Subsequent calls without the embeddings parameter will reutilize the faiss index built on the first call.
            >>> thresholds = [0.75] * len(weak_labels.rules)
            >>> weak_labels.extend_matrix(thresholds)
            >>> weak_labels.matrix()
        """
        abstains, supports = self._extend_matrix_preprocess()

        if embeddings is not None:
            self._extension_queries = self._find_dists_and_nearest(
                np.copy(embeddings).astype(np.float32), abstains, supports, gpu=gpu
            )
        elif self._extension_queries is None:
            raise ValueError("Embeddings are not optional the first time a matrix is extended.")
        dists, nearest = self._extension_queries

        self._extended_matrix = np.copy(self._matrix)
        new_points = [(dists[i] > thresholds[i]) for i in range(self._matrix.shape[1])]
        for i in range(self._matrix.shape[1]):
            self._extended_matrix[abstains[i][new_points[i]], i] = self._matrix[supports[i], i][
                nearest[i][new_points[i]]
            ]

        self._extend_matrix_postprocess()
Extends the weak label matrix through embeddings according to the similarity thresholds for each rule.

        Implementation based on `Epoxy <https://github.com/HazyResearch/epoxy>`__.

        Args:
            thresholds: An array of thresholds between 0.0 and 1.0, one for each column of the weak labels matrix.
                Each one stands for the minimum cosine similarity between two sentences for a rule to be extended.
            embeddings: Embeddings for each row of the weak label matrix.
                If not provided, we will use the ones from the last ``WeakLabels.extend_matrix()`` call.
            gpu: If True, perform FAISS similarity queries on GPU.

        Examples:
            >>> # Choose any model to generate the embeddings.
            >>> from sentence_transformers import SentenceTransformer
            >>> model = SentenceTransformer('all-mpnet-base-v2', device='cuda')
            >>>
            >>> # Generate the embeddings and set the thresholds.
            >>> weak_labels = {class_name}(dataset="my_dataset")
            >>> embeddings = np.array([ model.encode(rec.text) for rec in weak_labels.records() ])
            >>> thresholds = [0.6] * len(weak_labels.rules)
            >>>
            >>> # Extend the weak labels matrix.
            >>> weak_labels.extend_matrix(thresholds, embeddings)
            >>>
            >>> # Calling the method below will now retrieve the extended matrix.
            >>> weak_labels.matrix()
            >>>
            >>> # Subsequent calls without the embeddings parameter will reutilize the faiss index built on the first call.
            >>> thresholds = [0.75] * len(weak_labels.rules)
            >>> weak_labels.extend_matrix(thresholds)
            >>> weak_labels.matrix()
        
argilla-v1/src/argilla_v1/labeling/text_classification/weak_labels.py