使用文本嵌入的最佳便携方式是使用 Parquet 和 Polars

主要观点:现代文本嵌入(由大型语言模型生成)是生成式 AI 热潮中最有用的应用之一,作者创建了 32254 张 Magic: the Gathering 卡牌的文本嵌入。可以通过嵌入找到卡牌间的数学相似性,还能创建 2D UMAP 投影来识别有趣模式。生成嵌入后可使用向量数据库,但它们使用不便,其实在数据不大时无需向量数据库,利用 numpy 可快速计算最近邻相似度。存储嵌入有多种方式,如 CSV、.txt、pickle 等都不理想,Parquet 文件是较好的选择,能同时存储元数据和嵌入,且 Python 中的 pyarrow 和 polars 库可用于处理 Parquet 文件,polars 在处理嵌套数据方面表现更好,速度也更快,对于小型/非商业项目,Parquet 文件和 polars 的组合能满足需求,若数据量过大可考虑使用支持向量嵌入的老式数据库如 SQLite。

关键信息

  • 有 32254 张 Magic 卡牌的嵌入,可用于各种相似性分析。
  • numpy 可快速计算嵌入的相似度,平均 1.08ms 可计算 32254 个点积。
  • 存储嵌入的不良方式如 CSV 文本格式数据量大、.txt 保存加载慢且文件大、pickle 有安全风险等。
  • Parquet 文件是处理列数据的格式,可存储元数据和嵌入,pyarrow 和 polars 可用于处理。
  • polars 在处理嵌套数据和速度方面优于 pandas,能更方便地进行相似性计算和数据操作。

重要细节

  • 用 gte-modernbert-base 嵌入模型生成 Magic 卡牌嵌入,可在 Hugging Face 数据集获取。
  • 展示了 Wrath of God 卡牌及其相似卡牌,UMAP 投影可将卡牌聚类。
  • 介绍了各种存储嵌入方式的优缺点及示例代码。
  • 提供了交互式 Colab Notebook 可进行 Magic 卡牌相似性计算和过滤。
  • 提及在数据量过大时可考虑使用支持向量嵌入的 SQLite 数据库。
  • 代码处理 Magic 卡牌数据等相关代码在 GitHub 仓库。
阅读 10
0 条评论