-
Notifications
You must be signed in to change notification settings - Fork 237
Expand file tree
/
Copy pathsample.py
More file actions
39 lines (28 loc) · 1.56 KB
/
sample.py
File metadata and controls
39 lines (28 loc) · 1.56 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import random
from typing import Optional, TYPE_CHECKING
if TYPE_CHECKING: # pragma: no cover
from docarray.array.document import DocumentArray
class SampleMixin:
"""A mixin that provides search functionality to DocumentArrays"""
def sample(self, k: int, seed: Optional[int] = None) -> 'DocumentArray':
"""random sample k elements from :class:`DocumentArray` without replacement.
:param k: Number of elements to sample from the document array.
:param seed: initialize the random number generator, by default is None. If set will
save the state of the random function to produce certain outputs.
:return: A sampled list of :class:`Document` represented as :class:`DocumentArray`.
"""
if seed is not None:
random.seed(seed)
# NOTE, this could simplified to random.sample(self, k)
# without getting indices and itemgetter etc.
# however it's only work on DocumentArray.
sampled = random.sample(self, k)
from docarray.array.document import DocumentArray
return DocumentArray(sampled)
def shuffle(self, seed: Optional[int] = None) -> 'DocumentArray':
"""Randomly shuffle documents within the :class:`DocumentArray`.
:param seed: initialize the random number generator, by default is None. If set will
save the state of the random function to produce certain outputs.
:return: The shuffled list of :class:`Document` represented as :class:`DocumentArray`.
"""
return self.sample(len(self), seed=seed)