ITADN

Configurable truncation when rendering text

#53Closedshoyer 创建于 2025-01-09
feature-request
S
shoyercommented
Hierarchical models can create really large trees of objects. Treescope has nice control for shrinking reprs with `using_expansion_strategy(max_height=X)`, but this only works well for interactive HTML and displays that never wrap text (e.g., Google Colab). On the command line, long lines wrap around and make reprs unreadable. It would be really nice if Treescope has some built-in support for truncating such really long lines, perhaps by replacing objects with truncated reprs like `<dict>` or `{...}`. I would use this for implementing `__repr__` methods, e.g., for neural network modules. To reproduce: ```python from __future__ import annotations import dataclasses import numpy as np import treescope @dataclasses.dataclass class Config: batch_size: int = 128 num_features: int = 64 height: int = 1024 width: int = 1024 model_name: str = 'magic' nested_config: Config | None = None nested = { f'x{i}': { 'foo': np.random.randn(100, 100), 'bar': {'x': 1111, 'y': 2222, 'z': 3333, 'c': Config(nested_config=Config())}, 'baz': [3] * 7, } for i in range(5) } ``` Default rendering -- fine on GitHub, but unreadable when pasted into a console: ```python with treescope.using_expansion_strategy(max_height=10): print(treescope.render_to_text(nested)) ``` ``` { 'x0': {'foo': <np.ndarray float64(100, 100) ≈-0.0016 ±1.0 [≥-3.9, ≤3.9] nonzero:10_000>, 'bar': {'x': 1111, 'y': 2222, 'z': 3333, 'c': Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=None))}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, 'x1': {'foo': <np.ndarray float64(100, 100) ≈0.0019 ±1.0 [≥-3.5, ≤3.8] nonzero:10_000>, 'bar': {'x': 1111, 'y': 2222, 'z': 3333, 'c': Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=None))}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, 'x2': {'foo': <np.ndarray float64(100, 100) ≈0.01 ±1.0 [≥-3.5, ≤3.9] nonzero:10_000>, 'bar': {'x': 1111, 'y': 2222, 'z': 3333, 'c': Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=None))}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, 'x3': {'foo': <np.ndarray float64(100, 100) ≈0.003 ±0.99 [≥-3.8, ≤3.8] nonzero:10_000>, 'bar': {'x': 1111, 'y': 2222, 'z': 3333, 'c': Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=None))}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, 'x4': {'foo': <np.ndarray float64(100, 100) ≈0.022 ±1.0 [≥-3.7, ≤3.7] nonzero:10_000>, 'bar': {'x': 1111, 'y': 2222, 'z': 3333, 'c': Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=Config(batch_size=128, num_features=64, height=1024, width=1024, model_name='magic', nested_config=None))}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, } ``` One very naive attempt at truncation: ```python with treescope.using_expansion_strategy(max_height=10): print(treescope.render_to_text(nested)) def truncate_lines(text): return '\n'.join(t[:77] + '...' if len(t) > 80 else t for t in text.split('\n')) ``` ``` { 'x0': {'foo': <np.ndarray float64(100, 100) ≈-0.0016 ±1.0 [≥-3.9, ≤3.9] non... 'x1': {'foo': <np.ndarray float64(100, 100) ≈0.0019 ±1.0 [≥-3.5, ≤3.8] nonz... 'x2': {'foo': <np.ndarray float64(100, 100) ≈0.01 ±1.0 [≥-3.5, ≤3.9] nonzer... 'x3': {'foo': <np.ndarray float64(100, 100) ≈0.003 ±0.99 [≥-3.8, ≤3.8] nonz... 'x4': {'foo': <np.ndarray float64(100, 100) ≈0.022 ±1.0 [≥-3.7, ≤3.7] nonze... } ``` Example of what a better truncated repr might look like, with the tree truncated at a consistent depth: ``` { 'x0': {'foo': <np.ndarray>, 'bar': {...}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, 'x1': {'foo': <np.ndarray>, 'bar': {...}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, 'x2': {'foo': <np.ndarray>, 'bar': {...}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, 'x3': {'foo': <np.ndarray>, 'bar': {...}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, 'x4': {'foo': <np.ndarray>, 'bar': {...}, 'baz': [3, 3, 3, 3, 3, 3, 3]}, } ```
关闭于 2025-02-17 5 条评论