docarray/docarray/data/torch_dataset.py at main · docarray/docarray · GitHub

History

161 lines (118 loc) · 4.73 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

from typing import Callable, Dict, Generic, List, Optional, Type, TypeVar

from torch.utils.data import Dataset

from docarray import BaseDoc, DocList, DocVec

from docarray.typing import TorchTensor

from docarray.utils._internal._typing import change_cls_name, safe_issubclass

T_doc = TypeVar('T_doc', bound=BaseDoc)

class MultiModalDataset(Dataset, Generic[T_doc]):

"""

A dataset that can be used inside a PyTorch DataLoader.

In other words, it implements the PyTorch Dataset interface.

The preprocessing dictionary passed to the constructor consists of keys that are

field names and values that are functions that take a single argument and return

a single argument.

---

```python

from torch.utils.data import DataLoader

from docarray import DocList

from docarray.data import MultiModalDataset

from docarray.documents import TextDoc

def prepend_number(text: str):

return f"Number {text}"

docs = DocList[TextDoc](TextDoc(text=str(i)) for i in range(16))

ds = MultiModalDataset[TextDoc](docs, preprocessing={'text': prepend_number})

loader = DataLoader(ds, batch_size=4, collate_fn=MultiModalDataset[TextDoc].collate_fn)

for batch in loader:

print(batch.text)

```

---

Nested fields can be accessed by using dot notation.

The document itself can be accessed using the empty string as the key.

Transformations that operate on reference types (such as Documents) can optionally

not return a value.

The transformations will be applied according to their order in the dictionary.

---

```python

import torch

from torch.utils.data import DataLoader

from docarray import DocList, BaseDoc

from docarray.data import MultiModalDataset

from docarray.documents import TextDoc

class Thesis(BaseDoc):

title: TextDoc

class Student(BaseDoc):

thesis: Thesis

def embed_title(title: TextDoc):

title.embedding = torch.ones(4)

def normalize_embedding(thesis: Thesis):

thesis.title.embedding = thesis.title.embedding / thesis.title.embedding.norm()

def add_nonsense(student: Student):

student.thesis.title.embedding = student.thesis.title.embedding + int(

student.thesis.title.text

)

docs = DocList[Student](Student(thesis=Thesis(title=str(i))) for i in range(16))

ds = MultiModalDataset[Student](

docs,

preprocessing={

"thesis.title": embed_title,

"thesis": normalize_embedding,

"": add_nonsense,

)

loader = DataLoader(ds, batch_size=4, collate_fn=ds.collate_fn)

for batch in loader:

print(batch.thesis.title.embedding)

```

---

:param docs: the `DocList` to be used as the dataset

:param preprocessing: a dictionary of field names and preprocessing functions

"""

doc_type: Optional[Type[BaseDoc]] = None

__typed_ds__: Dict[Type[BaseDoc], Type['MultiModalDataset']] = {}

def __init__(

self, docs: 'DocList[T_doc]', preprocessing: Dict[str, Callable]

) -> None:

self.docs = docs

self._preprocessing = preprocessing

def __len__(self):

return len(self.docs)

def __getitem__(self, item: int):

doc = self.docs[item].copy(deep=True)

for field, preprocess in self._preprocessing.items():

if len(field) == 0:

doc = preprocess(doc) or doc

else:

acc_path = field.split('.')

_field_ref = doc

for attr in acc_path[:-1]:

_field_ref = getattr(_field_ref, attr)

attr = acc_path[-1]

value = getattr(_field_ref, attr)

setattr(_field_ref, attr, preprocess(value) or value)

return doc

@classmethod

def collate_fn(cls, batch: List[T_doc]):

doc_type = cls.doc_type

if doc_type:

batch_da = DocVec[doc_type]( # type: ignore

batch,

tensor_type=TorchTensor,

)

else:

batch_da = DocVec(batch, tensor_type=TorchTensor)

return batch_da

@classmethod

def __class_getitem__(cls, item: Type[BaseDoc]) -> Type['MultiModalDataset']:

if not safe_issubclass(item, BaseDoc):

raise ValueError(

f'{cls.__name__}[item] item should be a Document not a {item} '

)

if item not in cls.__typed_ds__:

global _TypedDataset

class _TypedDataset(cls): # type: ignore

doc_type = item

change_cls_name(

_TypedDataset, f'{cls.__name__}[{item.__name__}]', globals()

)

cls.__typed_ds__[item] = _TypedDataset

return cls.__typed_ds__[item]

docarray/docarray/data/torch_dataset.py at main · docarray/docarray · GitHub

Uh oh!

Expand file tree

Breadcrumbs

torch_dataset.py

Latest commit

History

Breadcrumbs

torch_dataset.py

File metadata and controls