使用云 AI 提供商创建嵌入

Cypher 函数 ai.text.embed 和过程 ai.text.embedBatch 允许您通过外部 AI 提供商为一个或多个文本片段生成嵌入。您需要拥有受支持提供商(OpenAI、Vertex AI、Azure OpenAI、Amazon Bedrock)中的任意一个的 API 令牌。

本页假设您已 导入了推荐数据集 并 配置好环境,展示如何基于电影的标题和情节生成并存储 Movie 节点的嵌入。

嵌入始终在 Neo4j 之外 生成,但 存储在 Neo4j 数据库中。

设置环境

编码函数是 Neo4j GenAI 插件的一部分。

  • 在 Aura 实例上,插件默认已启用,因此如果您使用的是 Aura 上的 Neo4j,无需执行其他操作。

  • 对于自托管实例,需要先安装插件。您可以将 neo4j‑genai.jar 文件从 /products 移动到 Neo4j 主目录下的 /plugins,或在启动 Docker 容器时添加额外参数 --env NEO4J_PLUGINS='["genai"]'。
    更多信息,请参见 配置 → 插件。

为电影创建嵌入

下面的示例从数据库中获取所有 Movie 节点,生成电影标题与情节拼接后的嵌入,并将其作为额外的 embedding 属性添加到每个节点上。

import neo4j


URI = '<database-uri>'
AUTH = ('<username>', '<password>')
DB_NAME = '<database-name>'  # examples: 'recommendations-5.26', 'neo4j'

openAI_token = '<OpenAI API token>'


def main():
    with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:  (1)
        driver.verify_connectivity()

        batch_size = 100
        batch_n = 1
        movies_batch = []
        with driver.session(database=DB_NAME) as session:
            # Fetch `Movie` nodes
            result = session.run('MATCH (m:Movie) RETURN m.plot AS plot, m.title AS title')
            for record in result:
                title = record.get('title')
                plot = record.get('plot')

                if title is not None and plot is not None:
                    movies_batch.append({
                        'title': title,
                        'plot': plot,
                        'to_encode': f'Title: {title}\nPlot: {plot}'  (2)
                    })

                # Import a batch; flush buffer
                if len(movies_batch) == batch_size:  (3)
                    import_batch(driver, movies_batch, batch_n)
                    movies_batch = []
                    batch_n += 1

            # Flush last batch
            import_batch(driver, movies_batch, batch_n)

        # Import complete, show counters
        records, _, _ = driver.execute_query('''
        MATCH (m:Movie WHERE m.embedding IS NOT NULL)
        RETURN count(*) AS countMoviesWithEmbeddings, size(m.embedding) AS embeddingSize
        ''', database_=DB_NAME)
        print(f"""
    Embeddings generated and attached to nodes.
    Movie nodes with embeddings: {records[0].get('countMoviesWithEmbeddings')}.
    Embedding size: {records[0].get('embeddingSize')}.
        """)


def import_batch(driver, nodes, batch_n):
    # Generate and store embeddings for Movie nodes
    driver.execute_query('''
    CALL ai.text.embedBatch($listToEncode, 'OpenAI', { (4)
      token: $token, model: 'text-embedding-3-small'
    }) YIELD index, vector  (5)
    MATCH (m:Movie {title: $movies[index].title, plot: $movies[index].plot})
    CALL db.create.setNodeVectorProperty(m, 'embedding', toFloatList(vector))  (6)
    ''', movies=nodes, listToEncode=[movie['to_encode'] for movie in nodes], token=openAI_token,
    database_=DB_NAME)
    print(f'Processed batch {batch_n}')


if __name__ == '__main__':
    main()

'''
Movie nodes with embeddings: 9083.
Embedding size: 1536.
'''
1 driver 对象是与您的 Neo4j 实例交互的接口。更多信息,请参见 使用 Neo4j 与 Python 构建应用程序。
2 需要让 OpenAI 编码为嵌入的字符串。
3 在提交完整批次到数据库之前,会先收集一定数量的嵌入。这可以避免一次性将整个数据集加载到内存中,从而防止潜在的超时(对大数据集尤为重要)。
4 过程 ai.text.embedBatch() 将批次提交给 OpenAI 进行编码。OpenAI 的模型 text-embedding-3-small 将文本嵌入为 1536 维向量(即包含 1536 个数字的列表)。有关受支持提供商和选项的列表,请参见 GenAI 提供商。始终对同一数据集使用相同的模型生成嵌入:选择一个模型并在整个项目中坚持使用。
5 从 ai.text.embedBatch 返回的 index 用于将嵌入与电影对应,从而可以检索每个电影节点并为其附加相应的嵌入;vector 的类型为 VECTOR,该存储方式仅在 Enterprise Edition 中可用。
6 导入查询在每个节点 m 上设置一个新的 embedding 属性,值为嵌入 vector。Cypher 过程 db.create.setNodeVectorProperty 比将向量存为列表更高效。由于 VECTOR 类型的属性仅在 Enterprise Edition 中可用,额外调用 toFloatList 可确保在所有版本间兼容。若需在关系上设置向量属性,请使用 db.create.setRelationshipVectorProperty。

使用 Enterprise Edition 时,您可以省略对 db.create.setNodeVectorProperty 的调用,直接通过 Cypher 子句 SET 将嵌入作为属性写入。相应的查询将会是如下形式:

CALL ai.text.embedBatch($listToEncode, 'OpenAI', {
  token: $token, model: 'text-embedding-3-small'
}) YIELD index, vector
MATCH (m:Movie {title: $movies[index].title, plot: $movies[index].plot})
SET m.embedding = vector

嵌入写入数据库后,您即可利用它们 比较电影之间的相似度。

© . This site is unofficial and not affiliated with Neo4j, Inc.