如何为大型 Pandas 数据框生成嵌入？

Question

我在为数据集生成嵌入时遇到问题。它由大约 16 000 000 条 reddit 评论组成（它们的正文 + 一些可以忽略不计的元数据）。我将它们存储为 CSV 文件，并使用 pd.read_csv() 从中生成 Pandas 数据框。

当我尝试生成嵌入时，困难的部分就开始了。我的程序的上一次迭代使用的数据集小了 3 个数量级，因此为每个评论生成嵌入是一项微不足道的任务 - 我只需为每个评论单独执行此操作，并将结果添加为新列。

然而，

df["embedding"] = self.embedder.embed_str(temp_df["author"])

方法已被证明不足以满足我的新数据集。我花了 6 个多小时等待它完成处理，结果又回到终端底部的

Killed

，这似乎是由于内存使用过多。我还尝试了并行批处理方法，但这导致进程的终止速度更快。

是否有一些更有效的方法可以做到这一点，或者我应该放弃并将嵌入留在训练过程中，而不是使其成为我的数据集的一部分？我将不胜感激有关此事的任何一般指导，因为这是我第一次接触数据科学。

为了提供额外的上下文，上述

self.embedder.embed_str()

方法如下：

    def embed_str(self, data: str) -> torch.Tensor:
        """Generates an embedding for the given str data.

        Args:
            data: The data to be embedded.

        Returns:
            A PyTorch tensor containing the embedding.
        """
        return self.model.encode(data)

，self.model 是一个 Jina Embeddings v2 SentenceTransformer，初始化如下：

self.model = SentenceTransformer(
            "jinaai/jina-embeddings-v2-base-en",
            trust_remote_code=True,
        )

为了了解更多背景信息，我在 Intel i5 10400F 和 16 GB RAM 上运行代码。

Answer 1

要为大型 Pandas 数据框生成嵌入，请考虑对文本数据使用 Word2Vec 或 Doc2Vec 等技术，或对数字数据使用 PCA 或 t-SNE 等方法。将数据帧分解为可管理的块以提高处理效率，并利用 Gensim 或 Scikit-learn 等库进行实施。