<commit_before>import io
from typing import Any

import numpy as np
from PIL import Image
from ipywidgets import widgets

from puzzle.puzzlepedia import _common
from puzzle.steps import step


def DebugDataWidget(s: step.Step) -> widgets.Widget:
  # TODO: Make data fetching lazy?
  try:
    data = s.get_debug_data()
  except NotImplementedError:
    data = '[no data]'
  return _data_widget(data)


def _data_widget(data: Any) -> widgets.Widget:
  if isinstance(data, np.ndarray) and data.dtype == np.uint8:
    f = io.BytesIO()
    Image.fromarray(data).save(f, 'png')
    return widgets.Image(value=f.getvalue())
  # Fallback assumes data is text.
  return widgets.HTML(_common.preformat_html(str(data)))
<commit_msg>Explicitly set width and height.<commit_after>import io
from typing import Any

import numpy as np
from PIL import Image
from ipywidgets import widgets

from puzzle.puzzlepedia import _common
from puzzle.steps import step


def DebugDataWidget(s: step.Step) -> widgets.Widget:
  # TODO: Make data fetching lazy?
  try:
    data = s.get_debug_data()
  except NotImplementedError:
    data = '[no data]'
  return _data_widget(data)


def _data_widget(data: Any) -> widgets.Widget:
  if isinstance(data, np.ndarray) and data.dtype == np.uint8:
    width, height = data.shape[:2]
    f = io.BytesIO()
    Image.fromarray(data).save(f, 'png')
    return widgets.Image(
        value=f.getvalue(),
        width=width,
        height=height,
    )
  # Fallback assumes data is text.
  return widgets.HTML(_common.preformat_html(str(data)))

{'fraction_non_alphanumeric': 0.10418006430868167, 'fraction_numerical': 0.0019292604501607716, 'mean_word_length': 3.908517350157729, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2129728', 'n_tokens_mistral': 518, 'n_tokens_neox': 509, 'n_words': 146}