def _find_dists_and_nearest(
        self,
        embeddings: np.ndarray,
        abstains: List[np.ndarray],
        support: List[np.ndarray],
        gpu: bool,
    ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
        """Helper method to extend the weak labels."""
        try:
            import faiss
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                "'faiss' must be installed to extend a weak label matrix! "
                "You can install 'faiss' with the commands: `pip install faiss-cpu` or `pip install faiss-gpu`"
            )
        faiss.normalize_L2(embeddings)
        embeddings_length = embeddings.shape[1]

        label_fn_indexes = [faiss.IndexFlatIP(embeddings_length) for i in range(self._matrix.shape[1])]

        if gpu:
            res = faiss.StandardGpuResources()
            label_fn_indexes = [faiss.index_cpu_to_gpu(res, 0, x) for x in label_fn_indexes]

        for i in range(self._matrix.shape[1]):
            label_fn_indexes[i].add(embeddings[support[i]])

        dists_and_nearest = []
        for i in tqdm(range(self._matrix.shape[1]), total=self._matrix.shape[1]):
            embs_query = np.copy(embeddings[abstains[i]])
            faiss.normalize_L2(embs_query)
            dists_and_nearest.append(label_fn_indexes[i].search(embs_query, 1))

        dists = [dist_and_nearest[0].flatten() for dist_and_nearest in dists_and_nearest]
        nearest = [dist_and_nearest[1].flatten() for dist_and_nearest in dists_and_nearest]

        return dists, nearest
Helper method to extend the weak labels.
argilla-v1/src/argilla_v1/labeling/text_classification/weak_labels.py