fix(document): fix embedding display with rich

hanxiao · hanxiao · commit 611b3910f350 · 2022-03-31T15:47:43.000+02:00
diff --git a/docarray/document/mixins/plot.py b/docarray/document/mixins/plot.py
@@ -17,9 +17,16 @@ def __rich_console__(self, console, options):
         my_table = Table('Attribute', 'Value', width=80, box=box.ROUNDED)
         for f in self.non_empty_fields:
             if f in ('embedding', 'tensor'):
-                from .rich_embedding import ColorBoxEmbedding
+                from ...math.ndarray import to_numpy_array
 
-                my_table.add_row(f, ColorBoxEmbedding(getattr(self, f)))
+                v = to_numpy_array(getattr(self, f))
+                if v.squeeze().ndim == 1:
+                    from .rich_embedding import ColorBoxEmbedding
+
+                    v = ColorBoxEmbedding(v.squeeze())
+                else:
+                    v = str(getattr(self, f))
+                my_table.add_row(f, v)
             elif f not in ('id', 'chunks', 'matches'):
                 my_table.add_row(f, str(getattr(self, f)))
         if my_table.rows:
diff --git a/docarray/document/mixins/rich_embedding.py b/docarray/document/mixins/rich_embedding.py
@@ -13,7 +13,7 @@
 
 class ColorBoxEmbedding:
     def __init__(self, array):
-        self._array = minmax_normalize(to_numpy_array(array).reshape(-1), (0, 5))
+        self._array = minmax_normalize(array, (0, 5))
 
     def __rich_console__(
         self, console: Console, options: ConsoleOptions
