test: undo changes in docker compose

Joan Fontanals Martinez · Joan Fontanals Martinez · commit bab809960e57 · 2023-07-11T16:30:45.000+02:00
Signed-off-by: Joan Fontanals Martinez &lt;joan.martinez@jina.ai&gt;
diff --git a/docarray/array/list_advance_indexing.py b/docarray/array/list_advance_indexing.py
@@ -1,5 +1,4 @@
 from typing import (
-    TYPE_CHECKING,
     Any,
     Iterable,
     List,
@@ -14,7 +13,18 @@
 import numpy as np
 from typing_extensions import SupportsIndex
 
-from docarray.utils._internal.misc import import_library
+from docarray.utils._internal.misc import (
+    is_torch_available,
+    is_tf_available,
+)
+
+torch_available = is_torch_available()
+if torch_available:
+    import torch
+tf_available = is_tf_available()
+if tf_available:
+    import tensorflow as tf  # type: ignore
+    from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor
 
 T_item = TypeVar('T_item')
 T = TypeVar('T', bound='ListAdvancedIndexing')
@@ -75,17 +85,20 @@ def _normalize_index_item(
             return item.tolist()
 
         # torch index types
-        if TYPE_CHECKING:
-            import torch
-        else:
-            torch = import_library('torch', raise_error=True)
-
-        allowed_torch_dtypes = [
-            torch.bool,
-            torch.int64,
-        ]
-        if isinstance(item, torch.Tensor) and (item.dtype in allowed_torch_dtypes):
-            return item.tolist()
+        if torch_available:
+
+            allowed_torch_dtypes = [
+                torch.bool,
+                torch.int64,
+            ]
+            if isinstance(item, torch.Tensor) and (item.dtype in allowed_torch_dtypes):
+                return item.tolist()
+
+        if tf_available:
+            if isinstance(item, tf.Tensor):
+                return item.numpy().tolist()
+            if isinstance(item, TensorFlowTensor):
+                return item.tensor.numpy().tolist()
 
         return item
 
diff --git a/tests/index/in_memory/test_in_memory.py b/tests/index/in_memory/test_in_memory.py
@@ -127,7 +127,7 @@ def test_with_text_doc_ndarray():
         [TextDoc(text='hey', embedding=np.random.rand(128)) for _ in range(200)]
     )
     index.index(docs)
-    res = index.find_batched(docs[0:10], search_field='embedding')
+    res = index.find_batched(docs[0:10], search_field='embedding', limit=5)
     assert len(res.documents) == 10
     for r in res.documents:
         assert len(r) == 5
@@ -144,7 +144,7 @@ def test_with_text_doc_tensorflow():
         ]
     )
     index.index(docs)
-    res = index.find_batched(docs[0:10], search_field='embedding')
+    res = index.find_batched(docs[0:10], search_field='embedding', limit=5)
     assert len(res.documents) == 10
     for r in res.documents:
         assert len(r) == 5
@@ -159,7 +159,7 @@ def test_with_text_doc_torch():
         [TextDoc(text='hey', embedding=torch.rand(128)) for _ in range(200)]
     )
     index.index(docs)
-    res = index.find_batched(docs[0:10], search_field='embedding')
+    res = index.find_batched(docs[0:10], search_field='embedding', limit=5)
     assert len(res.documents) == 10
     for r in res.documents:
         assert len(r) == 5
diff --git a/tests/index/weaviate/docker-compose.yml b/tests/index/weaviate/docker-compose.yml
@@ -1,4 +1,4 @@
-version: '3.3'
+version: '3.8'
 
 services:
   
@@ -24,4 +24,4 @@ services:
     LOG_LEVEL: debug # verbose
     LOG_FORMAT: text
     # LOG_LEVEL: trace # very verbose
-    GODEBUG: gctrace=1 # make go garbage collector verbose
+    GODEBUG: gctrace=1 # make go garbage collector verbose
