def write_prompt_files(
    dataset_name: str,
    dataset: Sequence[WinoDictExample],
    aux_dataset: Sequence[WinoDictExample],
    shots: int,
    strategy: Callable[[WinoDictExample], Prompts],
    seed: int,
):
  """Creates in-context learning given a strategy to write the prompts."""
  with tf.io.gfile.GFile(f'{dataset_name}_inputs', 'w') as f_inputs:
    with tf.io.gfile.GFile(f'{dataset_name}_labels', 'w') as f_outputs:
      for example_index, example in enumerate(dataset):
        candidates = [
            e for e in aux_dataset if _examples_compatible(e, example)
        ]
        if len(candidates) < shots:
          raise ValueError(
              f'Only {len(candidates)} examples are compatible with {example}')

        samples = random.Random(f'{example.get_id()}:{seed}').sample(
            candidates, shots)
        # Corresponds to the first correct example, the second wrong example
        # then second correct example and the second wrong example
        prompts = _build_few_shot_prompts(example, samples, strategy)
        indexes = (0, 1, 0, 1)
        if isinstance(example.answer1.idx, int):
          example_ids = [example.answer1.idx] * 2 + [example.answer2.idx] * 2
        else:
          example_ids = [2 * example_index] * 2 + [2 * example_index + 1] * 2

        for (inputs, targets), index, ids in zip(prompts, indexes, example_ids):
          f_inputs.write(f'{inputs.encode()}\t{targets.encode()}\n')
          f_outputs.write(f'(({ids}, {index}), {index == 0}, 1.0)\n')
Creates in-context learning given a strategy to write the prompts.
language/wino_dict/utils.py