Function load_lit_format

lit_nlp/api/dataset.py:420–441

Load data from LIT jsonl format.

(
    path: str, *args, id_fn=input_hash, **kw
)

def load_lit_format(
    path: str, *args, id_fn=input_hash, **kw
) -> Union[Dataset, IndexedDataset]:
  """Load data from LIT jsonl format."""
  with open(path + LIT_SPEC_EXTENSION, 'r') as fd:
    spec = serialize.from_json(fd.read())

  with open(path, 'r') as fd:
    examples = [serialize.from_json(line) for line in fd.readlines()]

  first_example_keys = set(ex.keys() if (ex := examples[0]) else [])
  # TODO(b/294233896): remove this once input representations are consolidated.
  if first_example_keys.issuperset({'id', 'data'}):
    return IndexedDataset(
        spec=spec,
        indexed_examples=cast(list[lit_types.IndexedInput], examples),
        id_fn=id_fn,
        *args,
        **kw,
    )
  else:
    return Dataset(spec=spec, examples=examples, *args, **kw)


# TODO(b/202210900): Remove "NoneDataset" once the LIT front-end constructs its