- Notifications
You must be signed in to change notification settings - Fork234
Represent, send, store and search multimodal data
License
docarray/docarray
Folders and files
| Name | Name | Last commit message | Last commit date | |
|---|---|---|---|---|
Repository files navigation
The data structure for multimodal data
⬆️DocArray v2: This readme refer to the second version of DocArray (starting at 0.30). If you want to use the oldDocArray v1 version (below 0.30) check out thedocarray-v1-fixe branch
DocArray is a library forrepresenting, sending and storing multi-modal data, perfect forMachine Learning applications.
Those are the three pillars of DocArray, and you can check them out individually:
DocArray handles your data while integrating seamlessly with the rest of yourPython and ML ecosystem:
- 🔥 DocArray has native compatibility forNumPy,PyTorch andTensorFlow, including formodel training use cases
- ⚡ DocArray is built onPydantic and out-of-the-box compatible withFastAPI andJina
- 📦 DocArray can index data in vector databases such asWeaviate,Qdrant,ElasticSearch as well asHNSWLib
- ⛓️ DocArray data can be sent as JSON overHTTP or asProtobuf overgRPC
💡Where are you coming from? Depending on your use case and background, there are different was to "get" DocArray.You can navigate to the following section for an explanation that should fit your mindest:
DocArray was released under the open-sourceApache License 2.0 in January 2022. It is currently a sandbox project underLF AI & Data Foundation.
DocArray allows you torepresent your data, in an ML-native way.
This is useful for different use cases:
- 🏃♀️ You aretraining a model, there are myriads of tensors of different shapes and sizes flying around, representing differentthings, and you want to keep a straight head about them
- ☁️ You areserving a model, for example through FastAPI, and you want to specify your API endpoints
- 🗂️ You areparsing data for later use in your ML or DS applications
💡Coming from Pydantic? If you're currently using Pydantic for the use cases above, you should be happy to hearthat DocArray is built on top of, and fully compatible with, Pydantic!Also, we havededicated section just for you!
Put simply, DocArray lets you represent your data in a dataclass-like way, with ML as a first class citizen:
fromdocarrayimportBaseDocfromdocarray.typingimportTorchTensor,ImageUrlimporttorch# Define your data modelclassMyDocument(BaseDoc):description:strimage_url:ImageUrl# could also be VideoUrl, AudioUrl, etc.image_tensor:TorchTensor[1704,2272,3]# you can express tensor shapes!# Stack multiple documents in a Document VectorfromdocarrayimportDocVecvec=DocVec[MyDocument]( [MyDocument(description="A cat",image_url="https://example.com/cat.jpg",image_tensor=torch.rand(1704,2272,3), ), ]*10)print(vec.image_tensor.shape)# (10, 1704, 2272, 3)
Click for more details
So let's take a closer look at how you can represent your data with DocArray:
fromdocarrayimportBaseDocfromdocarray.typingimportTorchTensor,ImageUrlfromtypingimportOptionalimporttorch# Define your data modelclassMyDocument(BaseDoc):description:strimage_url:ImageUrl# could also be VideoUrl, AudioUrl, etc.image_tensor:Optional[TorchTensor[1704,2272,3] ]# could also be NdArray or TensorflowTensorembedding:Optional[TorchTensor]
So not only can you define the types of your data, you can evenspecify the shape of your tensors!
Once you have your model in form of aDocument, you can work with it!
# Create a documentdoc=MyDocument(description="This is a photo of a mountain",image_url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",)# Load image tensor from URLdoc.image_tensor=doc.image_url.load()# Compute embedding with any model of your choicedefclip_image_encoder(image_tensor:TorchTensor)->TorchTensor:# dummy functionreturntorch.rand(512)doc.embedding=clip_image_encoder(doc.image_tensor)print(doc.embedding.shape)# torch.Size([512])
Of course you can compose Documents into a nested structure:
fromdocarrayimportBaseDocfromdocarray.documentsimportImageDoc,TextDocimportnumpyasnpclassMultiModalDocument(BaseDoc):image_doc:ImageDoctext_doc:TextDocdoc=MultiModalDocument(image_doc=ImageDoc(tensor=np.zeros((3,224,224))),text_doc=TextDoc(text='hi!'))
Of course, you rarely work with a single data point at a time, especially in Machine Learning applications.
That's why you can easily collect multipleDocuments:
When building or interacting with an ML system, usually you want to process multiple Documents (data points) at once.
DocArray offers two data structures for this:
DocVec: A vector ofDocuments. All tensors in theDocumentsare stacked up into a single tensor.Perfect for batch processing and use inside of ML models.DocList: A list ofDocuments. All tensors in theDocumentsare kept as-is.Perfect for streaming, re-ranking, and shuffling of data.
Let's take a look at them, starting withDocVec:
fromdocarrayimportDocVec,BaseDocfromdocarray.typingimportAnyTensor,ImageUrlimportnumpyasnpclassImage(BaseDoc):url:ImageUrltensor:AnyTensor# this allows torch, numpy, and tensor flow tensorsvec=DocVec[Image](# the DocVec is parametrized by your personal schema! [Image(url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",tensor=np.zeros((3,224,224)), )for_inrange(100) ])
As you can see in the code snippet above,DocVec isparametrized by the type of Document you want to use with it:DocVec[Image].
This may look slightly weird at first, but we're confident that you'll get used to it quickly!Besides, it allows us to do some cool things, like giving youbulk access to the fields that you defined in yourDocument:
tensor=vec.tensor# gets all the tensors in the DocVecprint(tensor.shape)# which are stacked up into a single tensor!print(vec.url)# you can bulk access any other field, too
The second data structure,DocList, works in a similar way:
fromdocarrayimportDocListdl=DocList[Image](# the DocList is parametrized by your personal schema! [Image(url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",tensor=np.zeros((3,224,224)), )for_inrange(100) ])
You can still bulk access the fields of yourDocument:
tensors=dl.tensor# gets all the tensors in the DocListprint(type(tensors))# as a list of tensorsprint(dl.url)# you can bulk access any other field, too
And you can insert, remove, and appendDocuments to yourDocList:
# appenddl.append(Image(url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",tensor=np.zeros((3,224,224)), ))# deletedeldl[0]# insertdl.insert(0,Image(url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",tensor=np.zeros((3,224,224)), ),)
And you can seamlessly switch betweenDocVec andDocList:
vec_2=dl.to_doc_vec()assertisinstance(vec_2,DocVec)dl_2=vec_2.to_doc_list()assertisinstance(dl_2,DocList)
DocArray allows you tosend your data, in an ML-native way.
This means there is native support forProtobuf and gRPC, on top ofHTTP and serialization to JSON, JSONSchema, Base64, and Bytes.
This is useful for different use cases:
- ☁️ You areserving a model, for example throughJina orFastAPI
- 🕸️ Youdistribute your model across machines and need to send your data between nodes
- ⚙️ You are building amicroservice architecture and need to send your data between microservices
💡Coming from FastAPI? If you're currently using FastAPI for the use cases above, you should be happy to hearthat DocArray is fully compatible with FastAPI!Also, we havededicated section just for you!
Whenever you want to send your data you need to serialize it, so let's take a look at how that works with DocArray:
fromdocarrayimportBaseDocfromdocarray.typingimportImageTorchTensorimporttorch# model your dataclassMyDocument(BaseDoc):description:strimage:ImageTorchTensor[3,224,224]# create a Documentdoc=MyDocument(description="This is a description",image=torch.zeros((3,224,224)),)# serialize it!proto=doc.to_protobuf()bytes_=doc.to_bytes()json=doc.json()# deserialize it!doc_2=MyDocument.from_protobuf(proto)doc_4=MyDocument.from_bytes(bytes_)doc_5=MyDocument.parse_raw(json)
Of course, serialization is not all you need.So check out how DocArray integrates with FatAPI and Jina.
Once you've modelled your data, and maybe sent it around, usually you want tostore it somewhere.But fret not! DocArray has you covered!
Document Stores let you, well, store your Documents, locally or remotely, all with the same user interface:
- 💿On disk as a file in your local file system
- 🪣 OnAWS S3
- ☁️ OnJina AI Cloud
See Document Store usage
The Document Store interface lets you push and pull Documents to and from multiple data sources, all with the same user interface.
As an example, let's take a look at how that would work with AWS S3 storage:
fromdocarrayimportDocListfromdocarray.documentsimportImageDocimportnumpyasnpdl=DocList[ImageDoc]( [ImageDoc(url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",tensor=np.zeros((3,224,224)), )for_inrange(100) ])# push the DocList to S3dl.push('s3://my-bucket/my-documents',show_progress=True)# pull the DocList from S3dl_2=DocList[ImageDoc].pull('s3://my-bucket/my-documents',show_progress=True)
Document Indexes let you index your Documents into avector database, for efficient similarity-based retrieval.
This is useful for:
- 🗨️ AugmentingLLMs and Chatbots with domain knowledge (Retrieval Augmented Generation)
- 🔍Neural search applications
- 💡Recommender systems
Currently, DocArray Document Indexes supportWeaviate,Qdrant,ElasticSearch, andHNSWLib, with more to come!
See Document Index usage
The Document Index interface lets you index and retrieve Documents from multiple vector databases, all with the same user interface.
It supports ANN vector search, text search, filtering, and hybrid search.
fromdocarrayimportDocList,BaseDocfromdocarray.indeximportHnswDocumentIndeximportnumpyasnpfromdocarray.typingimportImageUrl,ImageTensor,NdArrayclassImageDoc(BaseDoc):url:ImageUrltensor:ImageTensorembedding:NdArray[128]# create some datadl=DocList[ImageDoc]( [ImageDoc(url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",tensor=np.zeros((3,224,224)),embedding=np.random.random((128,)), )for_inrange(100) ])# create a Document Indexindex=HnswDocumentIndex[ImageDoc](work_dir='/tmp/test_index')# index your dataindex.index(dl)# find similar Documentsquery=dl[0]results,scores=index.find(query,limit=10,search_field='embedding')
Depending on your background and use case, there are different ways for you toget DocArray.Choose your own adventure!
Click to expand
If you are using DocArray v<0.30.0, you will be familiar with itsdataclass API.
DocArray v2 is that idea, taken seriously. EveryDocument is created through dataclass-like interface,courtesy ofPydantic.
This gives the following advantages:
- Flexibility: No need to conform to a fixed set of fields -- your data defines the schema
- Multi-modality: Easily store multiple modalities and multiple embeddings in the same Document
- Language agnostic: At its core, Documents are just dictionaries. This makes it easy to create and send them from any language, not just Python.
You may also be familiar with our old Document Stores for vector DB integration.They are now calledDocument Indexes and offer the following improvements (seehere for the new API):
- Hybrid search: You can now combine vector search with text search, and even filter by arbitrary fields
- Production-ready: The new Document Indexes are a much thinner wrapper around the various vector DB libraries, making them more robust and easier to maintain
- Increased flexibility: We strive to support any configuration or setting that you could perform through the DB's first-party client
For now, Document Indexes supportWeaviate,Qdrant,ElasticSearch, andHNSWLib, with more to come.
Click to expand
If you come from Pydantic, you can see DocArray Documents as juiced up Pydantic models, and DocArray as a collection of goodies around them.
More specifically, we set out tomake Pydantic fit for the ML world - not by replacing it, but by building on top of it!
This means that you get the following benefits:
- ML focused types: Tensor, TorchTensor, Embedding, ..., includingtensor shape validation
- Full compatibility withFastAPI
- DocList andDocVec generalize the idea of a model to asequence orbatch of models. Perfect foruse in ML models and other batch processing tasks.
- Types that are alive: ImageUrl can
.load()a URL to image tensor, TextUrl can load and tokenize text documents, etc. - Cloud-ready: Serialization toProtobuf for use with microservices andgRPC
- Pre-built multi-modal Documents for different data modalities: Image, Text, 3DMesh, Video, Audio and more. Note that all of these are valid Pydantic models!
- Document Stores andDocument Indexes let you store your data and retrieve it usingvector search
The most obvious advantage here isfirst-class support for ML centric data, such as {Torch, TF, ...}Tensor, Embedding, etc.
This includes handy features such as validating the shape of a tensor:
fromdocarrayimportBaseDocfromdocarray.typingimportTorchTensorimporttorchclassMyDoc(BaseDoc):tensor:TorchTensor[3,224,224]doc=MyDoc(tensor=torch.zeros(3,224,224))# worksdoc=MyDoc(tensor=torch.zeros(224,224,3))# works by reshapingtry:doc=MyDoc(tensor=torch.zeros(224))# fails validationexceptExceptionase:print(e)# tensor# Cannot reshape tensor of shape (224,) to shape (3, 224, 224) (type=value_error)classImage(BaseDoc):tensor:TorchTensor[3,'x','x']Image(tensor=torch.zeros(3,224,224))# workstry:Image(tensor=torch.zeros(3,64,128) )# fails validation because second dimension does not match thirdexceptExceptionase:print()try:Image(tensor=torch.zeros(4,224,224) )# fails validation because of the first dimensionexceptExceptionase:print(e)# Tensor shape mismatch. Expected(3, 'x', 'x'), got(4, 224, 224)(type=value_error)try:Image(tensor=torch.zeros(3,64) )# fails validation because it does not have enough dimensionsexceptExceptionase:print(e)# Tensor shape mismatch. Expected (3, 'x', 'x'), got (3, 64) (type=value_error)
Click to expand
If you come from PyTorch, you can see DocArray mainly as a way oforganizing your data as it flows through your model.
It offers you several advantages:
- Expresstensors shapes in type hints
- Group tensors that belong to the same object, e.g. an audio track and an image
- Go directly to deployment, by re-using your data model as aFastAPI orJina API schema
- Connect model components betweenmicroservices, using Protobuf and gRPC
DocArray can be used directly inside ML models to handle and represent multi-modal data.This allows you to reason about your data using DocArray's abstractions deep inside ofnn.Module,and provides a (FastAPI-compatible) schema that eases the transition between model training and model serving.
To see the effect of this, let's first observe a vanilla PyTorch implementation of a tri-modal ML model:
importtorchfromtorchimportnnimporttorchdefencoder(x):returntorch.rand(512)classMyMultiModalModel(nn.Module):def__init__(self):super().__init__()self.audio_encoder=encoder()self.image_encoder=encoder()self.text_encoder=encoder()defforward(self,text_1,text_2,image_1,image_2,audio_1,audio_2):embedding_text_1=self.text_encoder(text_1)embedding_text_2=self.text_encoder(text_2)embedding_image_1=self.image_encoder(image_1)embedding_image_2=self.image_encoder(image_2)embedding_audio_1=self.image_encoder(audio_1)embedding_audio_2=self.image_encoder(audio_2)return (embedding_text_1,embedding_text_2,embedding_image_1,embedding_image_2,embedding_audio_1,embedding_audio_2, )
Not very easy on the eyes if you ask us. And even worse, if you need to add one more modality you have to touch every part of your code base, changing theforward() return type and making a whole lot of changes downstream from that.
So, now let's see what the same code looks like with DocArray:
fromdocarrayimportDocList,BaseDocfromdocarray.documentsimportImageDoc,TextDoc,AudioDocfromdocarray.typingimportTorchTensorfromtorchimportnnimporttorchdefencoder(x):returntorch.rand(512)classPodcast(BaseDoc):text:TextDocimage:ImageDocaudio:AudioDocclassPairPodcast(BaseDoc):left:Podcastright:PodcastclassMyPodcastModel(nn.Module):def__init__(self):super().__init__()self.audio_encoder=encoder()self.image_encoder=encoder()self.text_encoder=encoder()defforward_podcast(self,docs:DocList[Podcast])->DocList[Podcast]:docs.audio.embedding=self.audio_encoder(docs.audio.tensor)docs.text.embedding=self.text_encoder(docs.text.tensor)docs.image.embedding=self.image_encoder(docs.image.tensor)returndocsdefforward(self,docs:DocList[PairPodcast])->DocList[PairPodcast]:docs.left=self.forward_podcast(docs.left)docs.right=self.forward_podcast(docs.right)returndocs
Looks much better, doesn't it?You instantly win in code readability and maintainability. And for the same price you can turn your PyTorch model into a FastAPI app and reuse your Documentschema definition (seebelow). Everything is handled in a pythonic manner by relying on type hints.
Click to expand
Similar to thePyTorch approach, you can also use DocArray with TensorFlow to handle and represent multi-modal data inside your ML model.
First off, to use DocArray with TensorFlow we first need to install it as follows:
pip install tensorflow==2.11.0pip install protobuf==3.19.0Compared to using DocArray with PyTorch, there is one main difference when using it with TensorFlow:
While DocArray'sTorchTensor is a subclass oftorch.Tensor, this is not the case for theTensorFlowTensor: Due to some technical limitations oftf.Tensor, DocArray'sTensorFlowTensor is not a subclass oftf.Tensor but rather stores atf.Tensor in its.tensor attribute.
How does this affect you? Whenever you want to access the tensor data to, let's say, do operations with it or hand it to your ML model, instead of handing over yourTensorFlowTensor instance, you need to access its.tensor attribute.
This would look like the following:
fromtypingimportOptionalfromdocarrayimportDocList,BaseDocimporttensorflowastfclassPodcast(BaseDoc):audio_tensor:Optional[AudioTensorFlowTensor]embedding:Optional[AudioTensorFlowTensor]classMyPodcastModel(tf.keras.Model):def__init__(self):super().__init__()self.audio_encoder=AudioEncoder()defcall(self,inputs:DocList[Podcast])->DocList[Podcast]:inputs.audio_tensor.embedding=self.audio_encoder(inputs.audio_tensor.tensor )# access audio_tensor's .tensor attributereturninputs
Click to expand
Documents are Pydantic Models (with a twist), and as such they are fully compatible with FastAPI!
But why should you use them, and not the Pydantic models you already know and love?Good question!
- Because of the ML-first features, types and validations,here
- Because DocArray can act as anORM for vector databases, similar to what SQLModel does for SQL databases
And to seal the deal, let us show you how easily Documents slot into your FastAPI app:
importnumpyasnpfromfastapiimportFastAPIfromhttpximportAsyncClientfromdocarrayimportBaseDocfromdocarray.documentsimportImageDocfromdocarray.typingimportNdArrayfromdocarray.base_docimportDocArrayResponseclassInputDoc(BaseDoc):img:ImageDocclassOutputDoc(BaseDoc):embedding_clip:NdArrayembedding_bert:NdArrayinput_doc=InputDoc(img=ImageDoc(tensor=np.zeros((3,224,224))))app=FastAPI()@app.post("/doc/",response_model=OutputDoc,response_class=DocArrayResponse)asyncdefcreate_item(doc:InputDoc)->OutputDoc:## call my fancy model to generate the embeddingsdoc=OutputDoc(embedding_clip=np.zeros((100,1)),embedding_bert=np.zeros((100,1)) )returndocasyncwithAsyncClient(app=app,base_url="http://test")asac:response=awaitac.post("/doc/",data=input_doc.json())resp_doc=awaitac.get("/docs")resp_redoc=awaitac.get("/redoc")
Just like a vanilla Pydantic model!
Click to expand
If you came across DocArray as a universal vector database client, you can best think of it asa new kind of ORM for vector databases.
DocArray's job is to take multi-modal, nested and domain-specific data and to map it to a vector database,store it there, and thus make it searchable:
fromdocarrayimportDocList,BaseDocfromdocarray.indeximportHnswDocumentIndeximportnumpyasnpfromdocarray.typingimportImageUrl,ImageTensor,NdArrayclassImageDoc(BaseDoc):url:ImageUrltensor:ImageTensorembedding:NdArray[128]# create some datadl=DocList[ImageDoc]( [ImageDoc(url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",tensor=np.zeros((3,224,224)),embedding=np.random.random((128,)), )for_inrange(100) ])# create a Document Indexindex=HnswDocumentIndex[ImageDoc](work_dir='/tmp/test_index2')# index your dataindex.index(dl)# find similar Documentsquery=dl[0]results,scores=index.find(query,limit=10,search_field='embedding')
Currently, DocArray supports the following vector databases:
- Weaviate
- Qdrant
- Elasticsearch v8 and v7
- HNSWlib as a local-first alternative
An integration ofOpenSearch is currently in progress.
Legacy versions of DocArray also supportRedis andMilvus, but these are not yet supported in the current version.
Of course this is only one thing that DocArray can do, so we encourage you to check out the rest of this readme!
To try out the alpha you can install it via git:
pip install"git+https://github.com/docarray/docarray@2023.01.18.alpha#egg=docarray[proto,torch,image]"...or from the latest development branch
pip install"git+https://github.com/docarray/docarray@feat-rewrite-v2#egg=docarray[proto,torch,image]"- Documentation
- Join our Discord server
- Donation to Linux Foundation AI&Data blog post
- "Legacy" DocArray github page
- "Legacy" DocArray documentation
DocArray is a trademark of LF AI Projects, LLC
About
Represent, send, store and search multimodal data
Topics
Resources
License
Code of conduct
Security policy
Uh oh!
There was an error while loading.Please reload this page.