def prepare_oasst_sft_results()

in ml/eval/generate.py [0:0]

• 14 lines of code
• 6 McCabe index (conditional complexity)

def prepare_oasst_sft_results(test_dataset, tokenizer, num_examples=None):
    """
    Prepare SFT results for a test dataset using a tokenizer.

    Parameters:
    - test_dataset: The dataset containing prompts and keys.
    - tokenizer: The tokenizer to process inputs and outputs.
    - num_examples: Optional; the number of examples to process.
                    If None, process the entire dataset.
    """
    sft_results = []
    # Limit dataset to num_examples if specified
    items = list(test_dataset.data.items())
    if num_examples is not None:
        items = items[:num_examples]

    for prompt, key in items:  # Iterate over limited dataset
        for i, j in key.pairs:  # Process each preference pair
            # Add prompt and corresponding chosen/rejected completions
            prompt = tokenizer.apply_chat_template(key.prompt, tokenize=False)
            output = key.generations[key.sft_index]

            # Keys are in alpacaeval format
            sft_results.append({
                "instruction": prompt,
                "output": output
            })
    return sft_results