Skip to main content
Nebius AI Studio 通过统一的接口提供对高质量嵌入模型的 API 访问权限。Nebius 嵌入模型将文本转换为捕捉语义含义的数值向量,使其在语义搜索、聚类和推荐等各种应用中非常有用。

概述

NebiusEmbeddings 类通过 LangChain 提供对 Nebius AI Studio 嵌入模型的访问权限。这些嵌入可用于语义搜索、文档相似度计算以及其他需要文本向量表示的自然语言处理(NLP)任务。

集成详情

  • 提供商:Nebius AI Studio
  • 模型类型:文本嵌入模型
  • 主要用例:生成文本的向量表示以用于语义相似度和检索
  • 可用模型:各种嵌入模型,包括 BAAI/bge-en-icl 等
  • 维度:因模型而异(通常为 1024-4096 维)

环境配置

安装

可通过 pip 安装 Nebius 集成包:
pip install -U langchain-nebius

凭据

Nebius 需要一个 API 密钥,该密钥可以作为初始化参数 api_key 传递,或设置为环境变量 NEBIUS_API_KEY。您可以通过在 Nebius AI Studio 上创建账户来获取 API 密钥。
import getpass
import os

# 确保您已将 API 密钥设置为环境变量
if "NEBIUS_API_KEY" not in os.environ:
    os.environ["NEBIUS_API_KEY"] = getpass.getpass("Enter your Nebius API key: ")

实例化

可以使用可选的 API 密钥和模型名称参数来实例化 NebiusEmbeddings 类:
from langchain_nebius import NebiusEmbeddings

# 初始化嵌入模型
embeddings = NebiusEmbeddings(
    # api_key="YOUR_API_KEY",  # 您可以直接传递 API 密钥
    model="BAAI/bge-en-icl"  # 默认嵌入模型
)

可用模型

支持的模型列表可在 studio.nebius.com/?modality=embedding 查看。

索引与检索

嵌入模型通常用于检索增强生成(RAG)流程中,既用于数据索引,也用于后续的检索。以下示例演示了如何将 NebiusEmbeddings 与向量存储结合使用以进行文档检索。
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document

# 准备文档
docs = [
    Document(
        page_content="Machine learning algorithms build mathematical models based on sample data"
    ),
    Document(page_content="Deep learning uses neural networks with many layers"),
    Document(page_content="Climate change is a major global environmental challenge"),
    Document(
        page_content="Neural networks are inspired by the human brain's structure"
    ),
]

# 创建向量存储
vector_store = FAISS.from_documents(docs, embeddings)

# 执行相似度搜索
query = "How does the brain influence AI?"
results = vector_store.similarity_search(query, k=2)

print("Search results for query:", query)
for i, doc in enumerate(results):
    print(f"Result {i + 1}: {doc.page_content}")
查询结果:大脑如何影响人工智能?
结果 1:神经网络受人类大脑结构的启发
结果 2:深度学习使用具有多层结构的神经网络

与 InMemoryVectorStore 配合使用

您还可以在轻量级应用中使用 InMemoryVectorStore
from langchain_core.vectorstores import InMemoryVectorStore

# 创建示例文本
text = "LangChain is a framework for developing applications powered by language models"

# 创建向量存储
vectorstore = InMemoryVectorStore.from_texts(
    [text],
    embedding=embeddings,
)

# 用作检索器
retriever = vectorstore.as_retriever()

# 检索相似文档
docs = retriever.invoke("What is LangChain?")
print(f"Retrieved document: {docs[0].page_content}")
检索到的文档:LangChain 是一个用于开发由语言模型驱动的应用程序的框架

直接使用

您可以直接使用 NebiusEmbeddings 类为文本生成嵌入,而无需使用向量存储。

嵌入单个文本

您可以使用 embed_query 方法嵌入单段文本:
query = "What is machine learning?"
query_embedding = embeddings.embed_query(query)

# 检查嵌入维度
print(f"Embedding dimension: {len(query_embedding)}")
print(f"First few values: {query_embedding[:5]}")
嵌入维度:4096
前几个值:[0.007419586181640625, 0.002246856689453125, 0.00193023681640625, -0.0066070556640625, -0.0179901123046875]

嵌入多个文本

您可以使用 embed_documents 方法一次性嵌入多个文本:
documents = [
    "Machine learning is a branch of artificial intelligence",
    "Deep learning is a subfield of machine learning",
    "Natural language processing deals with interactions between computers and human language",
]

document_embeddings = embeddings.embed_documents(documents)

# 检查结果
print(f"Number of document embeddings: {len(document_embeddings)}")
print(f"Each embedding has {len(document_embeddings[0])} dimensions")
文档嵌入数量:3
每个嵌入包含 4096 个维度

异步支持

NebiusEmbeddings 支持异步操作:
import asyncio


async def generate_embeddings_async():
    # 嵌入单个查询
    query_result = await embeddings.aembed_query("What is the capital of France?")
    print(f"Async query embedding dimension: {len(query_result)}")

    # 嵌入多个文档
    docs = [
        "Paris is the capital of France",
        "Berlin is the capital of Germany",
        "Rome is the capital of Italy",
    ]
    docs_result = await embeddings.aembed_documents(docs)
    print(f"Async document embeddings count: {len(docs_result)}")


await generate_embeddings_async()
异步查询嵌入维度:4096
异步文档嵌入数量:3

文档相似度示例

import numpy as np
from scipy.spatial.distance import cosine

# 创建一些文档
documents = [
    "Machine learning algorithms build mathematical models based on sample data",
    "Deep learning uses neural networks with many layers",
    "Climate change is a major global environmental challenge",
    "Neural networks are inspired by the human brain's structure",
]

# 嵌入文档
embeddings_list = embeddings.embed_documents(documents)


# 计算相似度的函数
def calculate_similarity(embedding1, embedding2):
    return 1 - cosine(embedding1, embedding2)


# 打印相似度矩阵
print("Document Similarity Matrix:")
for i, emb_i in enumerate(embeddings_list):
    similarities = []
    for j, emb_j in enumerate(embeddings_list):
        similarity = calculate_similarity(emb_i, emb_j)
        similarities.append(f"{similarity:.4f}")
    print(f"Document {i + 1}: {similarities}")
文档相似度矩阵:
文档 1: ['1.0000', '0.8282', '0.5811', '0.7985']
文档 2: ['0.8282', '1.0000', '0.5897', '0.8315']
文档 3: ['0.5811', '0.5897', '1.0000', '0.5918']
文档 4: ['0.7985', '0.8315', '0.5918', '1.0000']

API 参考

有关 Nebius AI Studio API 的更多详细信息,请访问 Nebius AI Studio 文档