Recherche sémantique avec Milvus et OpenAI

Ce guide montre comment l'API d'intégration d'OpenAI peut être utilisée avec la base de données vectorielle Milvus pour effectuer une recherche sémantique sur du texte.

Pour commencer

Avant de commencer, assurez-vous d'avoir la clé API OpenAI à portée de main ou d'en obtenir une sur le site Web d'OpenAI.

Les données utilisées dans cet exemple sont des titres de livres. Vous pouvez télécharger le jeu de données ici et le placer dans le même répertoire que celui où vous exécutez le code suivant.

Tout d'abord, installez le paquetage pour Milvus et OpenAI :

pip install --upgrade openai pymilvus

Nous sommes maintenant prêts à générer des embeddings et à utiliser la base de données vectorielle pour effectuer une recherche sémantique.

Recherche de titres de livres avec OpenAI et Milvus

Dans l'exemple suivant, nous chargeons les données relatives aux titres de livres à partir du fichier CSV téléchargé, nous utilisons le modèle d'intégration d'OpenAI pour générer des représentations vectorielles et nous les stockons dans la base de données vectorielles Milvus pour la recherche sémantique.

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"])

Avec toutes les données dans la base de données vectorielles Milvus, nous pouvons maintenant effectuer une recherche sémantique en générant une intégration vectorielle pour la requête et en effectuant une recherche vectorielle.

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")

Vous devriez obtenir le résultat suivant :

[
    {
        "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",
        },
    },
]