ベクトルデータベース「Milvus」を試す

!pip install -U pymilvus
!pip install openai

!pip freeze | egrep -i "milvus|openai"

milvus-lite==2.4.10
openai==1.47.1
pymilvus==2.4.6

from pymilvus import MilvusClient

milvus_client = MilvusClient("./milvus_demo.db")

DEBUG:pymilvus.milvus_client.milvus_client:Created new connection using: a0c04fbf5dc9448d9824a2db5cc03c33

$ ps auxw | grep -i milvus
root         600  0.1  0.2 194172 31876 ?        Sl   18:16   0:00 /usr/local/lib/python3.10/dist-packages/milvus_lite/lib/milvus /content/milvus_demo.db unix:/tmp/tmp2fz3z7tc_milvus_demo.db.sock ERROR /content/.milvus_demo.db.lock
root         716  0.0  0.0   6484  2280 pts/2    S+   18:16   0:00 grep --color=auto -i milvus

$ file milvus_demo.db 
milvus_demo.db: SQLite 3.x database, last written using SQLite version 3045000, file counter 1, database pages 2, cookie 0x1, schema 4, UTF-8, version-valid-for 1

milvus_client.create_collection(
    collection_name="demo_collection",
    dimension=1536,  
)

DEBUG:pymilvus.milvus_client.milvus_client:Successfully created collection: demo_collection
DEBUG:pymilvus.milvus_client.milvus_client:Successfully created an index on collection: demo_collection

import os
from openai import OpenAI
from google.colab import userdata

os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

oepnai_client = OpenAI()

docs = [
    "人工知能は、1956年に学問分野として確立されました。",
    "アラン・チューリングは、AIの研究を本格的に行った最初の人物です。",
    "ロンドン近郊の町、メイダ・ヴェールで生まれたチューリングは、イングランド南部で育ちました。",
]

response = oepnai_client.embeddings.create(
    input=docs,
    model = "text-embedding-3-small"
)

vectors = [data.embedding for data in response.data]

data = [ {"id": idx, "vector": vector, "text": docs[idx], "subject": "history"} for idx, vector in enumerate(vectors)]

res = milvus_client.insert(
    collection_name="demo_collection",
    data=data
)

import json

query = "人工知能は学問として確立したのはいつですか？"

response = oepnai_client.embeddings.create(
    input=[query],
    model = "text-embedding-3-small"
)

res = milvus_client.search(
    collection_name="demo_collection",
    data=[response.data[0].embedding],
    filter="subject == 'history'",
    limit=2,
    output_fields=["text", "subject"],
)
print(json.dumps(res, indent=4, ensure_ascii=False))

[
    [
        {
            "id": 0,
            "distance": 0.836178719997406,
            "entity": {
                "text": "人工知能は、1956年に学問分野として確立されました。",
                "subject": "history"
            }
        },
        {
            "id": 1,
            "distance": 0.3081776797771454,
            "entity": {
                "text": "アラン・チューリングは、AIの研究を本格的に行った最初の人物です。",
                "subject": "history"
            }
        }
    ]
]

res = milvus_client.search(
    collection_name="demo_collection",
    data=[response.data[0].embedding],
    filter="subject != 'history'",
    limit=2,
    output_fields=["text", "subject"],
)
print(json.dumps(res, indent=4, ensure_ascii=False))

[
    []
]

res = milvus_client.query(
    collection_name="demo_collection",
    filter="subject == 'history'",
    output_fields=["text", "subject"],
)
print(json.dumps(res, indent=4, ensure_ascii=False))

[
    {
        "id": 0,
        "text": "人工知能は、1956年に学問分野として確立されました。",
        "subject": "history"
    },
    {
        "id": 1,
        "text": "アラン・チューリングは、AIの研究を本格的に行った最初の人物です。",
        "subject": "history"
    },
    {
        "id": 2,
        "text": "ロンドン近郊の町、メイダ・ヴェールで生まれたチューリングは、イングランド南部で育ちました。",
        "subject": "history"
    }
]

res = milvus_client.delete(
    collection_name="demo_collection",
    filter="id == 2",
)
print(json.dumps(res, indent=4, ensure_ascii=False))

[
    2
]

res = milvus_client.query(
    collection_name="demo_collection",
    output_fields=["text", "subject"],
    limit=5,
)
print(json.dumps(res, indent=4, ensure_ascii=False))

[
    {
        "id": 0,
        "text": "人工知能は、1956年に学問分野として確立されました。",
        "subject": "history"
    },
    {
        "id": 1,
        "text": "アラン・チューリングは、AIの研究を本格的に行った最初の人物です。",
        "subject": "history"
    }
]