milvus-logo
LFAI
Home
  • Integrações

Pesquisa semântica com Milvus e OpenAI

Open In Colab GitHub Repository

Este guia mostra como a API de incorporação do OpenAI pode ser utilizada com a base de dados vetorial Milvus para efetuar uma pesquisa semântica de texto.

Introdução

Antes de começar, certifique-se de que tem a chave da API OpenAI pronta, ou obtenha uma no sítio Web da OpenAI.

Os dados utilizados neste exemplo são títulos de livros. Pode descarregar o conjunto de dados aqui e colocá-lo no mesmo diretório onde executa o código seguinte.

Primeiro, instale o pacote do Milvus e do OpenAI:

pip install --upgrade openai pymilvus

Se estiver a utilizar o Google Colab, para ativar as dependências que acabou de instalar, poderá ter de reiniciar o tempo de execução. (Clique no menu "Runtime" (Tempo de execução) na parte superior do ecrã e selecione "Restart session" (Reiniciar sessão) no menu pendente).

Com isso, estamos prontos para gerar embeddings e usar o banco de dados de vetores para realizar a pesquisa semântica.

Pesquisa de títulos de livros com o OpenAI e o Milvus

No exemplo seguinte, carregamos os dados dos títulos dos livros a partir do ficheiro CSV descarregado, utilizamos o modelo de incorporação OpenAI para gerar representações vectoriais e armazenamo-las na base de dados vetorial Milvus para pesquisa semântica.

from openai import OpenAI
from pymilvus import MilvusClient

MODEL_NAME = "text-embedding-3-small"  # Which model to use, please check https://platform.openai.com/docs/guides/embeddings for available models
DIMENSION = 1536  # Dimension of vector embedding

# Connect to OpenAI with API Key.
openai_client = OpenAI(api_key="<YOUR_OPENAI_API_KEY>")

docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

vectors = [
    vec.embedding
    for vec in openai_client.embeddings.create(input=docs, model=MODEL_NAME).data
]

# Prepare data to be stored in Milvus vector database.
# We can store the id, vector representation, raw text and labels such as "subject" in this case in Milvus.
data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(docs))
]


# Connect to Milvus, all data is stored in a local file named "milvus_openai_demo.db"
# in current directory. You can also connect to a remote Milvus server following this
# instruction: https://milvus.io/docs/install_standalone-docker.md.
milvus_client = MilvusClient(uri="milvus_openai_demo.db")
COLLECTION_NAME = "demo_collection"  # Milvus collection name
# Create a collection to store the vectors and text.
if milvus_client.has_collection(collection_name=COLLECTION_NAME):
    milvus_client.drop_collection(collection_name=COLLECTION_NAME)
milvus_client.create_collection(collection_name=COLLECTION_NAME, dimension=DIMENSION)

# Insert all data into Milvus vector database.
res = milvus_client.insert(collection_name="demo_collection", data=data)

print(res["insert_count"])

Quanto ao argumento de MilvusClient:

  • Definir o uri como um ficheiro local, por exemplo./milvus.db, é o método mais conveniente, uma vez que utiliza automaticamente o Milvus Lite para armazenar todos os dados neste ficheiro.
  • Se tiver uma grande escala de dados, pode configurar um servidor Milvus mais eficiente em docker ou kubernetes. Nesta configuração, utilize o uri do servidor, por exemplo,http://localhost:19530, como o seu uri.
  • Se pretender utilizar o Zilliz Cloud, o serviço de nuvem totalmente gerido para o Milvus, ajuste os endereços uri e token, que correspondem ao Public Endpoint e à chave Api no Zilliz Cloud.

Com todos os dados na base de dados de vectores do Milvus, podemos agora efetuar uma pesquisa semântica gerando uma incorporação de vectores para a consulta e realizar uma pesquisa de vectores.

queries = ["When was artificial intelligence founded?"]

query_vectors = [
    vec.embedding
    for vec in openai_client.embeddings.create(input=queries, model=MODEL_NAME).data
]

res = milvus_client.search(
    collection_name=COLLECTION_NAME,  # target collection
    data=query_vectors,  # query vectors
    limit=2,  # number of returned entities
    output_fields=["text", "subject"],  # specifies fields to be returned
)

for q in queries:
    print("Query:", q)
    for result in res:
        print(result)
    print("\n")

Deverá ver o seguinte como resultado:

[
    {
        "id": 0,
        "distance": -0.772376537322998,
        "entity": {
            "text": "Artificial intelligence was founded as an academic discipline in 1956.",
            "subject": "history",
        },
    },
    {
        "id": 1,
        "distance": -0.58596271276474,
        "entity": {
            "text": "Alan Turing was the first person to conduct substantial research in AI.",
            "subject": "history",
        },
    },
]

Traduzido porDeepLogo

Try Managed Milvus for Free

Zilliz Cloud is hassle-free, powered by Milvus and 10x faster.

Get Started
Feedback

Esta página foi útil?