|
4 | 4 | from collections import defaultdict |
5 | 5 | from dataclasses import dataclass, field |
6 | 6 | from typing import ( |
| 7 | + TYPE_CHECKING, |
7 | 8 | Any, |
8 | 9 | Dict, |
9 | 10 | Generator, |
|
21 | 22 | ) |
22 | 23 |
|
23 | 24 | import numpy as np |
24 | | -from elastic_transport import NodeConfig |
25 | | -from elasticsearch import Elasticsearch |
26 | | -from elasticsearch.helpers import parallel_bulk |
27 | 25 | from pydantic import parse_obj_as |
28 | 26 |
|
29 | 27 | import docarray.typing |
|
32 | 30 | from docarray.typing import AnyTensor |
33 | 31 | from docarray.typing.tensor.abstract_tensor import AbstractTensor |
34 | 32 | from docarray.typing.tensor.ndarray import NdArray |
35 | | -from docarray.utils._internal.misc import is_tf_available, is_torch_available |
| 33 | +from docarray.utils._internal.misc import import_library |
36 | 34 | from docarray.utils.find import _FindResult, _FindResultBatched |
37 | 35 |
|
38 | 36 | TSchema = TypeVar('TSchema', bound=BaseDoc) |
39 | 37 | T = TypeVar('T', bound='ElasticDocIndex') |
40 | 38 |
|
41 | 39 | ELASTIC_PY_VEC_TYPES: List[Any] = [list, tuple, np.ndarray, AbstractTensor] |
42 | 40 |
|
43 | | -if is_torch_available(): |
44 | | - import torch |
45 | 41 |
|
46 | | - ELASTIC_PY_VEC_TYPES.append(torch.Tensor) |
| 42 | +if TYPE_CHECKING: |
| 43 | + from elastic_transport import NodeConfig |
| 44 | + from elasticsearch import Elasticsearch |
| 45 | + from elasticsearch.helpers import parallel_bulk |
| 46 | +else: |
| 47 | + elasticsearch = import_library('elasticsearch', raise_error=True) |
| 48 | + from elasticsearch import Elasticsearch |
| 49 | + from elasticsearch.helpers import parallel_bulk |
| 50 | + |
| 51 | + elastic_transport = import_library('elastic_transport', raise_error=True) |
| 52 | + from elastic_transport import NodeConfig |
| 53 | + |
| 54 | + torch = import_library('torch', raise_error=False) |
| 55 | + tf = import_library('tensorflow', raise_error=False) |
47 | 56 |
|
48 | | -if is_tf_available(): |
49 | | - import tensorflow as tf # type: ignore |
50 | 57 |
|
| 58 | +if torch is not None: |
| 59 | + ELASTIC_PY_VEC_TYPES.append(torch.Tensor) |
| 60 | + |
| 61 | +if tf is not None: |
51 | 62 | from docarray.typing import TensorFlowTensor |
52 | 63 |
|
53 | 64 | ELASTIC_PY_VEC_TYPES.append(tf.Tensor) |
|
0 commit comments