DocArray — 适合多模态 AI 数据组织、搜索与服务编排

让我们仔细看看如何使用 DocArray 表示数据：

from docarray import BaseDoc
from docarray.typing import TorchTensor, ImageUrl
from typing import Optional
import torch


# 定义您的数据模型
class MyDocument(BaseDoc):
    description: str
    image_url: ImageUrl  # 也可以是 VideoUrl、AudioUrl 等
    image_tensor: Optional[
        TorchTensor[1704, 2272, 3]
    ] = None  # 也可以是 NdArray 或 TensorflowTensor
    embedding: Optional[TorchTensor] = None

因此，您不仅可以定义数据的类型，甚至可以指定张量的形状！

# 创建一个文档
doc = MyDocument(
    description="This is a photo of a mountain",
    image_url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",
)

# 从 URL 加载图像张量
doc.image_tensor = doc.image_url.load()


# 使用您选择的任何模型计算嵌入
def clip_image_encoder(image_tensor: TorchTensor) -> TorchTensor:  # 虚拟函数
    return torch.rand(512)


doc.embedding = clip_image_encoder(doc.image_tensor)

print(doc.embedding.shape)  # torch.Size([512])

组合嵌套文档

当然，您可以将文档组合成嵌套结构：

from docarray import BaseDoc
from docarray.documents import ImageDoc, TextDoc
import numpy as np


class MultiModalDocument(BaseDoc):
    image_doc: ImageDoc
    text_doc: TextDoc


doc = MultiModalDocument(
    image_doc=ImageDoc(tensor=np.zeros((3, 224, 224))), text_doc=TextDoc(text='hi!')
)

您很少一次只处理单个数据点，尤其是在机器学习应用中。因此，您可以轻松地收集多个 Document：

收集多个 Document

在构建或与 ML 系统交互时，通常希望一次处理多个 Document（数据点）。

DocArray 为此提供了两种数据结构：

• DocVec：一个 Document 的向量。文档中的所有张量被堆叠为单个张量。非常适合批量处理和在 ML 模型内部使用。
• DocList：一个 Document 的列表。文档中的所有张量都保持原样。非常适合数据流、重排序和洗牌。

让我们从 DocVec 开始看：

from docarray import DocVec, BaseDoc
from docarray.typing import AnyTensor, ImageUrl
import numpy as np


class Image(BaseDoc):
    url: ImageUrl
    tensor: AnyTensor  # 允许 torch、numpy 和 tensor flow 张量


vec = DocVec[Image](  # DocVec 由您的个人模式参数化！
    [
        Image(
            url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",
            tensor=np.zeros((3, 224, 224)),
        )
        for _ in range(100)
    ]
)

在上面的代码片段中，DocVec 由您希望与其一起使用的文档类型参数化：DocVec[Image]。

这可能起初看起来有些奇怪，但我们相信您会很快习惯！此外，它让我们可以做很酷的事情，比如对您在文档中定义的字段进行批量访问：

tensor = vec.tensor  # 获取 DocVec 中的所有张量
print(tensor.shape)  # 它们被堆叠成一个张量！
print(vec.url)  # 您也可以批量访问任何其他字段

第二种数据结构 DocList 的工作方式类似：

from docarray import DocList

dl = DocList[Image](  # DocList 由您的个人模式参数化！
    [
        Image(
            url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",
            tensor=np.zeros((3, 224, 224)),
        )
        for _ in range(100)
    ]
)

您仍然可以批量访问文档的字段：

tensors = dl.tensor  # 获取 DocList 中的所有张量
print(type(tensors))  # 作为张量列表
print(dl.url)  # 您也可以批量访问任何其他字段

您可以向 DocList 中插入、删除和追加文档：

# 追加
dl.append(
    Image(
        url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",
        tensor=np.zeros((3, 224, 224)),
    )
)
# 删除
del dl[0]
# 插入
dl.insert(
    0,
    Image(
        url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",
        tensor=np.zeros((3, 224, 224)),
    ),
)

您可以在 DocVec 和 DocList 之间无缝切换：

vec_2 = dl.to_doc_vec()
assert isinstance(vec_2, DocVec)

dl_2 = vec_2.to_doc_list()
assert isinstance(dl_2, DocList)

如果您使用的是 DocArray 0.30.0 或更低版本，您会熟悉其 dataclass API。

DocArray >=0.30 是那个理念的正式体现。 每个文档都通过类似 dataclass 的接口创建，这得益于 Pydantic。

这带来了以下优势：
- 灵活性： 无需遵守一组固定的字段——您的数据定义模式
- 多模态性： 文档本质上是字典。这使得从任何语言（不仅仅是 Python）创建和发送它们都变得简单。

您可能还熟悉我们旧的用于向量数据库集成的 Document Stores。它们现在被称为 Document Indexes，并提供以下改进（有关新 API，请参阅此处）：

• 混合搜索： 您现在可以将向量搜索与文本搜索结合，甚至可以按任意字段进行过滤
• 生产就绪： 新的 Document Indexes 是各种向量数据库库更薄的封装，使其更健壮且更易于维护
• 更高的灵活性： 我们努力支持您通过数据库的官方客户端可以执行的任何配置或设置

目前，Document Indexes 支持 Weaviate、Qdrant、ElasticSearch、Redis、Mongo Atlas、精确最近邻搜索和 HNSWLib，更多功能即将推出。

如果您从 Pydantic 转来，可以将 DocArray 文档视为增强版的 Pydantic 模型，而 DocArray 则是围绕它们的好用工具的集合。

具体来说，我们致力于让 Pydantic 适合 ML 世界——不是通过替换它，而是在其基础上构建！

这意味着您将获得以下好处：

• 面向 ML 的类型：Tensor、TorchTensor、Embedding 等，包括张量形状验证
• 与 FastAPI 完全兼容
• DocList 和 DocVec 将模型的概念推广到模型的_序列_或_批量_。非常适合在 ML 模型和批量处理任务中使用。
• 活跃类型：ImageUrl 可以 .load() 一个 URL 到图像张量，TextUrl 可以加载和分词文本文档等。
• 云就绪：序列化为 Protobuf 以便与微服务和 gRPC 一起使用
• 预构建的多模态文档，用于不同的数据模态：图像、文本、3D 网格、视频、音频等。注意所有这些都是有效的 Pydantic 模型！
• Document Stores 和 Document Indexes 允许您存储数据并使用向量搜索检索

这里最明显的优势是对 ML 中心数据（如 {Torch、TF，...}Tensor、Embedding 等）的一流支持。

这包括张量形状验证等便捷功能：

from docarray import BaseDoc
from docarray.typing import TorchTensor
import torch


class MyDoc(BaseDoc):
    tensor: TorchTensor[3, 224, 224]


doc = MyDoc(tensor=torch.zeros(3, 224, 224))  # 有效
doc = MyDoc(tensor=torch.zeros(224, 224, 3))  # 通过重塑有效

try:
    doc = MyDoc(tensor=torch.zeros(224))  # 验证失败
except Exception as e:
    print(e)
    # tensor
    # Cannot reshape tensor of shape (224,) to shape (3, 224, 224) (type=value_error)


class Image(BaseDoc):
    tensor: TorchTensor[3, 'x', 'x']


Image(tensor=torch.zeros(3, 224, 224))  # 有效

try:
    Image(
        tensor=torch.zeros(3, 64, 128)
    )  # 验证失败，因为第二维与第三维不匹配
except Exception as e:
    print()


try:
    Image(
        tensor=torch.zeros(4, 224, 224)
    )  # 验证失败，因为第一维
except Exception as e:
    print(e)
    # Tensor shape mismatch. Expected(3, 'x', 'x'), got(4, 224, 224)(type=value_error)

try:
    Image(
        tensor=torch.zeros(3, 64)
    )  # 验证失败，因为维度不足
except Exception as e:
    print(e)
    # Tensor shape mismatch. Expected (3, 'x', 'x'), got (3, 64) (type=value_error)

如果您从 PyTorch 转来，可以将 DocArray 主要视为在数据流经模型时组织数据的一种方式。

它为您提供了几个优势：

• 在类型提示中表达张量形状
• 将属于同一对象的张量分组，例如音频轨道和图像
• 通过复用您的数据模型作为 FastAPI 或 Jina API 的架构，直接进入部署阶段
• 使用 Protobuf 和 gRPC 连接微服务之间的模型组件

DocArray 可以直接在 ML 模型内部使用，以处理和表示多模态数据。这使您能够使用 DocArray 在 nn.Module 内部对数据进行抽象，并提供与 FastAPI 兼容的架构，从而简化从模型训练到模型服务的过渡。

为了了解其效果，首先看一个简单的三模态 ML 模型的纯 PyTorch 实现：

import torch
from torch import nn


def encoder(x):
    return torch.rand(512)


class MyMultiModalModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.audio_encoder = encoder()
        self.image_encoder = encoder()
        self.text_encoder = encoder()

    def forward(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,
        )

这不像我们说的那么易读。更糟糕的是，如果您需要添加一种模态，您必须触及代码库的每个部分，更改 forward() 的返回类型并在下游进行大量更改。

现在，让我们看看使用 DocArray 的相同代码：

from docarray import DocList, BaseDoc
from docarray.documents import ImageDoc, TextDoc, AudioDoc
from docarray.typing import TorchTensor
from torch import nn
import torch


def encoder(x):
    return torch.rand(512)


class Podcast(BaseDoc):
    text: TextDoc
    image: ImageDoc
    audio: AudioDoc


class PairPodcast(BaseDoc):
    left: Podcast
    right: Podcast


class MyPodcastModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.audio_encoder = encoder()
        self.image_encoder = encoder()
        self.text_encoder = encoder()

    def forward_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)

        return docs

    def forward(self, docs: DocList[PairPodcast]) -> DocList[PairPodcast]:
        docs.left = self.forward_podcast(docs.left)
        docs.right = self.forward_podcast(docs.right)

        return docs

看起来好多了，对吧？您在代码可读性和可维护性上立即获胜。同时，您可以将 PyTorch 模型变成 FastAPI 应用，并重用您的 Document模式定义（参见下方）。一切均以 Pythonic 的方式通过类型提示处理。

像 PyTorch 方法一样，您也可以将 DocArray 与 TensorFlow 结合使用，以处理和表示 ML 模型中的多模态数据。

首先，要使用 DocArray 与 TensorFlow，我们需要按如下方式安装它：

pip install tensorflow==2.12.0
pip install protobuf==3.19.0

与使用 PyTorch 相比，在使用 TensorFlow 时有一个主要区别：虽然 DocArray 的 TorchTensor 是 torch.Tensor 的子类，但 TensorFlowTensor 并非如此：由于 tf.Tensor 的一些技术限制，DocArray 的 TensorFlowTensor 不是 tf.Tensor 的子类，而是在其 .tensor 属性中存储了一个 tf.Tensor。

这对您有何影响？每当您想要访问张量数据以进行操作或将其传递给 ML 模型时，您需要访问其 .tensor 属性，而不是直接传递 TensorFlowTensor 实例。

看起来如下：

from typing import Optional

from docarray import DocList, BaseDoc

import tensorflow as tf


class Podcast(BaseDoc):
    audio_tensor: Optional[AudioTensorFlowTensor] = None
    embedding: Optional[AudioTensorFlowTensor] = None


class MyPodcastModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.audio_encoder = AudioEncoder()

    def call(self, inputs: DocList[Podcast]) -> DocList[Podcast]:
        inputs.audio_tensor.embedding = self.audio_encoder(
            inputs.audio_tensor.tensor
        )  # 访问 audio_tensor 的 .tensor 属性
        return inputs

Document 是 Pydantic 模型（略有不同），因此它们与 FastAPI 完全兼容！

但为什么要使用它们，而不是您已经熟悉和喜爱的 Pydantic 模型呢？好问题！

• 因为 ML 优先的特性、类型和验证，见此处
• 因为 DocArray 可以作为向量数据库的 ORM，类似于 SQLModel 对 SQL 数据库的作用

接下来，让我们展示 Document 如何轻松地融入您的 FastAPI 应用：

import numpy as np
from fastapi import FastAPI
from docarray.base_doc import DocArrayResponse
from docarray import BaseDoc
from docarray.documents import ImageDoc
from docarray.typing import NdArray, ImageTensor


class InputDoc(BaseDoc):
    img: ImageDoc
    text: str


class OutputDoc(BaseDoc):
    embedding_clip: NdArray
    embedding_bert: NdArray


app = FastAPI()


def model_img(img: ImageTensor) -> NdArray:
    return np.zeros((100, 1))


def model_text(text: str) -> NdArray:
    return np.zeros((100, 1))


@app.post("/embed/", response_model=OutputDoc, response_class=DocArrayResponse)
async def create_item(doc: InputDoc) -> OutputDoc:
    doc = OutputDoc(
        embedding_clip=model_img(doc.img.tensor), embedding_bert=model_text(doc.text)
    )
    return doc


input_doc = InputDoc(text='', img=ImageDoc(tensor=np.random.random((3, 224, 224))))

async with AsyncClient(app=app, base_url="http://test") as ac:
    response = await ac.post("/embed/", data=input_doc.json())

就像一个普通的 Pydantic 模型一样！

Jina 已采用 docarray 作为其表示和序列化 Document 的库。

Jina 允许您服务使用 DocArray 构建的模型和服务，让您可以充分利用 DocArray 的序列化能力来部署和扩展这些应用。

import numpy as np
from jina import Deployment, Executor, requests
from docarray import BaseDoc, DocList
from docarray.documents import ImageDoc
from docarray.typing import NdArray, ImageTensor


class InputDoc(BaseDoc):
    img: ImageDoc
    text: str


class OutputDoc(BaseDoc):
    embedding_clip: NdArray
    embedding_bert: NdArray


def model_img(img: ImageTensor) -> NdArray:
    return np.zeros((100, 1))


def model_text(text: str) -> NdArray:
    return np.zeros((100, 1))


class MyEmbeddingExecutor(Executor):
    @requests(on='/embed')
    def encode(self, docs: DocList[InputDoc], **kwargs) -> DocList[OutputDoc]:
        ret = DocList[OutputDoc]()
        for doc in docs:
            output = OutputDoc(
                embedding_clip=model_img(doc.img.tensor),
                embedding_bert=model_text(doc.text),
            )
            ret.append(output)
        return ret


with Deployment(
    protocols=['grpc', 'http'], ports=[12345, 12346], uses=MyEmbeddingExecutor
) as dep:
    resp = dep.post(
        on='/embed',
        inputs=DocList[InputDoc](
            [InputDoc(text='', img=ImageDoc(tensor=np.random.random((3, 224, 224))))]
        ),
        return_type=DocList[OutputDoc],
    )
    print(resp)

如果您是通过 DocArray 作为通用向量数据库客户端了解的，可以最好地将其视为向量数据库的一种新型 ORM。DocArray 的工作是将多模态、嵌套和领域特定的数据映射到向量数据库，存储在其中，从而使其可搜索：

from docarray import DocList, BaseDoc
from docarray.index import HnswDocumentIndex
import numpy as np

from docarray.typing import ImageUrl, ImageTensor, NdArray


class ImageDoc(BaseDoc):
    url: ImageUrl
    tensor: ImageTensor
    embedding: NdArray[128]


# 创建一些数据
dl = 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 _ in range(100)
    ]
)

# 创建一个文档索引
index = HnswDocumentIndex[ImageDoc](work_dir='/tmp/test_index2')


# 索引您的数据
index.index(dl)

# 查找相似的文档
query = dl[0]
results, scores = index.find(query, limit=10, search_field='embedding')

目前，DocArray 支持以下向量数据库：

• Weaviate
• Qdrant
• Elasticsearch v8 和 v7
• Redis
• Milvus
• ExactNNMemorySearch 作为本地替代方案，支持精确 kNN 搜索。
• HNSWlib 作为本地优先的 ANN 替代方案
• Mongo Atlas

OpenSearch 的集成正在开发中。

当然，这只是 DocArray 能做的其中一件事，所以我们鼓励您查看本说明的其余部分！

使用 DocArray，您可以通过 Langchain 将外部数据连接到 LLM。DocArray 让您可以自由地建立灵活的文模式架构，并从不同的后端选择存储文档。创建文档索引后，您可以使用 DocArrayRetriever 将其连接到您的 Langchain 应用。

通过以下命令安装 Langchain：

pip install langchain

1. 定义模式并创建文档：

from docarray import BaseDoc, DocList
from docarray.typing import NdArray
from langchain.embeddings.openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings()


# 定义文档模式
class MovieDoc(BaseDoc):
    title: str
    description: str
    year: int
    embedding: NdArray[1536]


movies = [
    {"title": "#1 title", "description": "#1 description", "year": 1999},
    {"title": "#2 title", "description": "#2 description", "year": 2001},
]

# 嵌入 `description` 并创建文档
docs = DocList[MovieDoc](
    MovieDoc(embedding=embeddings.embed_query(movie["description"]), **movie)
    for movie in movies
)

2. 使用任何支持的 backend 初始化文档索引：

from docarray.index import (
    InMemoryExactNNIndex,
    HnswDocumentIndex,
    WeaviateDocumentIndex,
    QdrantDocumentIndex,
    ElasticDocIndex,
    RedisDocumentIndex,
    MongoDBAtlasDocumentIndex,
)

# 选择合适的后端并初始化数据
db = InMemoryExactNNIndex[MovieDoc](docs)

3. 最后，初始化一个检索器并将其集成到您的链中！

from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain
from langchain.retrievers import DocArrayRetriever


# 创建检索器
retriever = DocArrayRetriever(
    index=db,
    embeddings=embeddings,
    search_field="embedding",
    content_field="description",
)

# 在您的链中使用该检索器
model = ChatOpenAI()
qa = ConversationalRetrievalChain.from_llm(model, retriever=retriever)

另外，您可以使用内置的向量存储。Langchain 支持两种向量存储： DocArrayInMemorySearch 和 DocArrayHnswSearch。两者都用户友好，最适合中小型数据集。