Configurable truncation when rendering text
feature-request
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 条评论