Using SuperDuperDB to Accelerate AI Development on MongoDB Atlas Vector Search

1
python -m pip install -U superduperdb[apis]

1
from superduperdb import superduper
2
3
db = superduper("mongodb+srv://<user>:<password>@...mongodb.net/documents")

1
export OPENAI_API_KEY="sk-..."

1
from superduperdb.ext.openai.model import OpenAIEmbedding
2
3
model = OpenAIEmbedding(
4
    identifier='text-embedding-ada-002', model='text-embedding-ada-002')

1
>>> model.predict('This is a test', one=True)
2
[-0.008146246895194054,
3
 -0.0036965329200029373,
4
 -0.0006024622125551105,
5
 -0.005724836140871048,
6
 -0.02455105632543564,
7
 0.01614714227616787,
8
...]

1
import sentence_transformers
2
from superduperdb.components.model import Model

1
from superduperdb import vector

1
model = Model(
2
    identifier='all-MiniLM-L6-v2',
3
    object=sentence_transformers.SentenceTransformer('all-MiniLM-L6-v2'),
4
    encoder=vector(shape=(384,)),
5
    predict_method='encode',
6
    postprocess=lambda x: x.tolist(),
7
    batch_predict=True,
8
)

1
>>> model.predict('This is a test', one=True)
2
[-0.008146246895194054,
3
 -0.0036965329200029373,
4
 -0.0006024622125551105,
5
 -0.005724836140871048,
6
 -0.02455105632543564,
7
 0.01614714227616787,
8
...]

1
curl -O https://superduperdb-public.s3.eu-west-1.amazonaws.com/pymongo.json

1
import json
2
from superduperdb.backends.mongodb.query import Collection
3
from superduperdb.base.document import Document as D
4
5
with open('pymongo.json') as f:
6
    data = json.load(f)
7
8
db.execute(
9
    Collection('documents').insert_many([D(r) for r in data])
10
)

1
>>> r = db.execute(Collection('documents').find_one())
2
>>> r
3
Document({
4
    'key': 'pymongo.mongo_client.MongoClient', 
5
    'parent': None, 
6
    'value': '\nClient for a MongoDB instance, a replica set, or a set of mongoses.\n\n', 
7
    'document': 'mongo_client.md',
8
    'res': 'pymongo.mongo_client.MongoClient',
9
    '_fold': 'train',
10
    '_id': ObjectId('652e460f6cc2a5f9cc21db4f')
11
})

1
>>> r.unpack()
2
{'key': 'pymongo.mongo_client.MongoClient',
3
 'parent': None,
4
 'value': '\nClient for a MongoDB instance, a replica set, or a set of mongoses.\n\n',
5
 'document': 'mongo_client.md',
6
 'res': 'pymongo.mongo_client.MongoClient',
7
 '_fold': 'train',
8
 '_id': ObjectId('652e460f6cc2a5f9cc21db4f')}

1
from superduperdb import VectorIndex
2
from superduperdb import Listener
3
4
db.add(
5
    VectorIndex(
6
        identifier='pymongo-docs',
7
        indexing_listener=Listener(
8
            model=model,
9
            key='value',
10
            select=Collection('documents').find(),
11
            predict_kwargs={'max_chunk_size': 1000},
12
        ),
13
    )
14
)

1
>>> db.load("model", 'all-MiniLM-L6-v2')

1
>>> db.execute(Collection('documents').find_one()).unpack()
2
{'key': 'pymongo.mongo_client.MongoClient',
3
 'parent': None,
4
 'value': '\nClient for a MongoDB instance, a replica set, or a set of mongoses.\n\n',
5
 'document': 'mongo_client.md',
6
 'res': 'pymongo.mongo_client.MongoClient',
7
 '_fold': 'train',
8
 '_id': ObjectId('652e460f6cc2a5f9cc21db4f'),
9
 '_outputs': {'value': {'text-embedding-ada-002': [-0.024740776047110558,
10
    0.013489063829183578,
11
    0.021334229037165642,
12
    -0.03423869237303734,
13
    ...]}}}

1
>>> query = 'Query the database'
2
>>> result = db.execute(
3
...    Collection('documents')
4
...        .like(D({'value': query}), vector_index='pymongo-docs', n=5)
5
...        .find({}, {'value': 1, 'key': 1})
6
... )
7
>>> for r in result:
8
...    print(r.unpack())
9
10
{'key': 'find', 'value': '\nQuery the database.\n\nThe filter argument is a query document that all results\nmust match. For example:\n\n`pycon\n>>> db'}
11
{'key': 'database_name', 'value': '\nThe name of the database this command was run against.\n\n'}
12
{'key': 'aggregate', 'value': '\nPerform a database-level aggregation.\n\nSee the [aggregation pipeline](https://mongodb.com/docs/manual/reference/operato'}
13
{'key': 'alive', 'value': '\nDoes this cursor have the potential to return more data?\n\n'}
14
{'key': 'pymongo.cursor.CursorType', 'value': '\n'}