Class IndexedDataset

lit_nlp/api/dataset.py:252–417

Dataset with additional indexing information.

class IndexedDataset(Dataset):
  """Dataset with additional indexing information."""

  _index: dict[ExampleId, IndexedInput] = {}

  def _normalize_example(
      self, data: JsonDict, ex_id: ExampleId, meta: lit_types.InputMetadata
  ):
    return types.MappingProxyType(dict(data, _id=ex_id, _meta=meta))

  def index_inputs(
      self, examples: list[lit_types.JsonDict]
  ) -> list[IndexedInput]:
    """Create indexed versions of inputs."""
    indexed = []
    for example in examples:
      ex_id = example.get(INPUT_ID_FIELD, self.id_fn(example))
      ex_meta = example.get(
          INPUT_META_FIELD,
          lit_types.InputMetadata(added=None, parentId=None, source=None),
      )
      indexed.append(
          IndexedInput(
              data=types.MappingProxyType(
                  example | {INPUT_ID_FIELD: ex_id, INPUT_META_FIELD: ex_meta}
              ),
              id=ex_id,
              meta=ex_meta,
          )
      )
    return indexed

  def __init__(
      self,
      *args,
      id_fn: Optional[IdFnType] = None,
      indexed_examples: Optional[list[IndexedInput]] = None,
      **kw,
  ):
    # The base Dataset class will initialize self._examples in this call to
    # super().__init__(), which may or may not include the _id and _meta fields.
    super().__init__(*args, **kw)
    self.id_fn = id_fn if id_fn is not None else input_hash

    if indexed_examples:
      self._indexed_examples = indexed_examples
      # Ensure that all indexed exampls provide a readonly view of their data.
      for ie in self._indexed_examples:
        if not isinstance((ie_data := ie['data']), types.MappingProxyType):
          ie['data'] = self._normalize_example(ie_data, ie['id'], ie['meta'])
    else:
      self._indexed_examples = self.index_inputs(self._examples)

    self._examples = [
        self._normalize_example(ex['data'], ex['id'], ex.get('meta', {}))
        for ex in self._indexed_examples
    ]
    self._index = {ex['id']: ex for ex in self._indexed_examples}