class composer.datasets.InContextLearningCodeEvalDataset(generations_per_sample, pass_at_k=1, *args, **kwargs)[source]#

A dataset that constructs batches for in-context learning code evaluation.

The input format is expected to be a jsonl file with the following fields:

  • task_id: Label of given task

  • prompt: The code snippet that must be completed

  • entry_point: The entry to the function/code snippet to generate

  • canonical_solution: Working solution

  • test: The checker code that will run to completion if the code generation is valid and otherwise throw assertion

  • test_inputs: List of test inputs

  • test_outputs: List of test outputs

  • language: The language of the code snippet

Each batch then consists of the following the structure

  • input_ids: Input tensor batch x seqlen x num tokens

  • mode: Indicates to the model that this is an ICL task and may rely on a custom code path to properly update metrics

  • mode: Always set to โ€˜generateโ€™

  • labels: Exact solution for the coding problem

  • prompts: Prompt for the task

  • entry_points: List of entry points

  • test_inputs: List of test inputs

  • test_outputs: List of test outputs

  • languages: List of languages

  • pass_at_k: Passed value for pass_at_k

  • generation_kwargs: Dictionary of kwargs neeeded for generation. Includes the following, which will be individually overwritten by keys in generaiton_kwargs if set (see for more details):

    • pad_token_id: ID for padding token, derived automatically

    • num_beams: How many beams to search for generations, set to 1

    • do_sample: Determines whether model is sampling or greedily decoding. Always set to True

    • use_cache: Whether or not to use past key values to speed up sampling. Always set to True

Additional Args:

generations_per_sample (int) (defaults to 1): The number of independently computed returned sequences for each element in the batch pass_at_k (int) (defaults to 1): k for how many chances the model gets to write passing code

tokenize_example(prompt_and_fewshot, ctxt, example)[source]#

Adds extra code task details to the example dictionary. See InContextLearningDataset for more details