fix: summary for document

anna-charlotte · anna-charlotte · commit 5d453d4829e7 · 2023-01-19T15:06:59.000+01:00
Signed-off-by: anna-charlotte &lt;charlotte.gerhaher@jina.ai&gt;
diff --git a/docarray/base_document/mixins/plot.py b/docarray/base_document/mixins/plot.py
@@ -1,9 +1,16 @@
+import colorsys
 from typing import Any, Optional, TypeVar
 
 import numpy as np
+from rich.color import Color
+from rich.console import Console, ConsoleOptions, RenderResult
+from rich.measure import Measurement
+from rich.segment import Segment
+from rich.style import Style
 from rich.tree import Tree
 
 import docarray
+from docarray.math.helper import minmax_normalize
 from docarray.typing import ID
 
 T = TypeVar('T', bound=Any)
@@ -37,7 +44,8 @@ def _plot_recursion(node: T, tree: Optional[Tree] = None) -> Tree:
                 for i, d in enumerate(value):
                     if i == 2:
                         PlotMixin._plot_recursion(
-                            f' ... {len(value) - 2} more Docs', _match_tree
+                            f' ... {len(value) - 2} more {d.__class__} documents',
+                            _match_tree,
                         )
                         break
                     PlotMixin._plot_recursion(d, _match_tree)
@@ -51,6 +59,7 @@ def __rich_console__(self, console, options):
         kls = self.__class__.__name__
         id_abbrv = getattr(self, 'id')[:7]
         yield f":page_facing_up: [b]{kls}" f"[/b]: [cyan]{id_abbrv} ...[cyan]"
+
         from collections.abc import Iterable
 
         import torch
@@ -64,7 +73,7 @@ def __rich_console__(self, console, options):
         for k, v in self.__dict__.items():
             col_1, col_2 = '', ''
 
-            if k.startswith('_') or isinstance(v, ID) or v is None:
+            if isinstance(v, ID) or k.startswith('_') or v is None:
                 continue
             elif isinstance(v, str):
                 col_1 = f'{k}: {v.__class__.__name__}'
@@ -73,7 +82,14 @@ def __rich_console__(self, console, options):
                     col_2 += f' ... (length: {len(v)})'
             elif isinstance(v, np.ndarray) or isinstance(v, torch.Tensor):
                 col_1 = f'{k}: {v.__class__.__name__}'
-                col_2 = f'{type(v)} in shape {v.shape}, dtype: {v.dtype}'
+
+                if isinstance(v, torch.Tensor):
+                    v = v.detach().cou().numpy()
+                if v.squeeze().ndim == 1 and len(v) < 1000:
+                    col_2 = ColorBoxArray(v.squeeze())
+                else:
+                    col_2 = f'{type(v)} of shape {v.shape}, dtype: {v.dtype}'
+
             elif isinstance(v, tuple) or isinstance(v, list):
                 col_1 = f'{k}: {v.__class__.__name__}'
                 for i, x in enumerate(v):
@@ -88,7 +104,31 @@ def __rich_console__(self, console, options):
             else:
                 continue
 
-            my_table.add_row(col_1, text.Text(col_2))
+            if not isinstance(col_2, ColorBoxArray):
+                col_2 = text.Text(col_2)
+            my_table.add_row(col_1, col_2)
 
         if my_table.rows:
             yield my_table
+
+
+class ColorBoxArray:
+    def __init__(self, array):
+        self._array = minmax_normalize(array, (0, 5))
+
+    def __rich_console__(
+        self, console: Console, options: ConsoleOptions
+    ) -> RenderResult:
+        h = 0.75
+        for idx, y in enumerate(self._array):
+            lightness = 0.1 + ((y / 5) * 0.7)
+            r, g, b = colorsys.hls_to_rgb(h, lightness + 0.7 / 10, 1.0)
+            color = Color.from_rgb(r * 255, g * 255, b * 255)
+            yield Segment('&#9604;', Style(color=color, bgcolor=color))
+            if idx != 0 and idx % options.max_width == 0:
+                yield Segment.line()
+
+    def __rich_measure__(
+        self, console: "Console", options: ConsoleOptions
+    ) -> Measurement:
+        return Measurement(1, options.max_width)
