import torch
from typing import Optional
from docarray import BaseDocument, DocumentArray
from docarray.typing import TorchTensor
from docarray.documents import Text as BaseText, Image


class Tokens(BaseDocument):
    input_ids: TorchTensor[48]
    attention_mask: TorchTensor


class Text(BaseText):
    tokens: Optional[Tokens]


class PairTextImage(BaseDocument):
    text: Text
    image: Image


def get_da(N: Optional[int] = None):
    da = DocumentArray[PairTextImage](
        PairTextImage(
            text=Text(
                text=str(i),
                tokens=Tokens(
                    input_ids=torch.randn(48), attention_mask=torch.randn(48)
                ),
            ),
            image=Image(url=f"Images/{i}.png"),
        )
        for i in range(N)
    )
    return da.stack()

da = get_da(10)
da.to_protobuf()