Example: generating and running code

Warning: This notebook runs LLM-generated code without any checks. Run at your own risk.

Loading a code model:

from guidance import models, gen
import guidance
import re

model = models.Transformers("Qwen/Qwen2.5-Coder-1.5B")

Loading the HumanEval dataset:

from datasets import load_dataset
dataset = load_dataset("openai_humaneval")

Let’s write a very simple baseline

import re
import guidance
@guidance
def baseline(lm, prompt):
    r = re.findall(r'def (.*?)\(', prompt)
    name = r[-1]
    lm += f'Here is an implementation of {name}:\n'
    lm += '```python\n' + prompt + gen(max_tokens=200, stop_regex=r'```|if __name__|def test|\n[^\s]', name='program')
    lm = lm.set('program', prompt + lm['program'])
    return lm

idx = 121
prompt = dataset['test']['prompt'][idx]
lm = model + baseline(prompt)
baseline_program = lm['program']

Here is simple function to evaluate a generated program with the HumanEval evaluation tests:

# Returns True if it passes the evaluation tests, False otherwise
def eval_program(program, i):
    namespace = {}
    # Loads the `check` function
    exec(dataset['test']['test'][i], namespace)
    try:
        # Executes the function definition
        exec(program, namespace)
    except Exception as e:
        print("Program not valid")
        print(e)
        return False
    name = dataset['test']['entry_point'][i]
    try:
        # Run the unit tests
        namespace['check'](namespace[name])
        # If we get here, we passed the tests
        return True
    except Exception as e:
        print("Program ran, but hit exception")
        print(e)
        return False

lm['program']

'\ndef solution(lst):\n    """Given a non-empty list of integers, return the sum of all of the odd elements that are in even positions.\n    \n\n    Examples\n    solution([5, 8, 7, 1]) ==> 12\n    solution([3, 3, 3, 3, 3]) ==> 9\n    solution([30, 13, 24, 321]) ==>0\n    """\n    # Initialize the sum to 0\n    sum = 0\n    \n    # Iterate over the list\n    for i in range(len(lst)):\n        # Check if the element is odd and in an even position\n        if lst[i] % 2 != 0 and i % 2 == 0:\n            # Add the element to the sum\n            sum += lst[i]\n    \n    # Return the sum\n    return sum\n'

eval_program(lm['program'], idx)

True

Let’s try another one:

idx = 71
prompt = dataset['test']['prompt'][idx]
lm = model + baseline(prompt)
baseline_program = lm['program']

eval_program(lm['program'], idx)

True

Let’s manually check the basic example from the docstring:

exec(lm['program'])
triangle_area(3, 4, 5) # should be 6.00

6.0

This suggests an improvement: let’s extract the tests on the docstrings, and only return a program if it passes at least those tests.

First, let’s write a simple prompt to extract the examples from the docstring into tests

@guidance
def write_tests(lm, prompt):
    r = re.findall(r'def (.*?)\(', prompt)
    name = r[-1]
    lm += '```python\n' + prompt + '    pass\n'
    lm += f'\ndef test_{name}():\n'
    lm += '    """Turns the example(s) in the docstring above into asserts"""\n'
    examples = re.findall(rf'{re.escape(name)}\((.*?)\)\s*(?:==|=>)\s*([^\n]+)', prompt)
    if not examples:
        raise ValueError(f'Could not find docstring examples for {name}')
    args = []
    expected = []
    for arg, result in examples:
        arg = arg.strip()
        result = result.strip()
        lm += f'   assert {name}({arg}) == {result}\n'
        args.append(arg)
        expected.append(result)
    lm = lm.set('args', args)
    lm = lm.set('expected', expected)
    return lm

lm = model + write_tests(prompt)
args = lm['args']
expected = lm['expected']

The LM went beyond extracting tests, it also generated a few of its own. While some of these may be incorrect, at least we have the original ones as well.

What’s more, we already stored the inputs and expected results in the lm object:

# (input, expected output)
list(zip(lm['args'], lm['expected']))

[('3, 4, 5', '6.00'), ('1, 2, 10', '-1')]

Let’s combine the baseline and the test generation prompts into a single guidance function:

@guidance
def reconstruct_tests(lm, name, args, expected):
    """Helper to format tests nicely"""
    lm += f'def test_{name}():\n'
    for arg, e in zip(args, expected):
        lm += f'   assert {name}({arg}) == {e}\n'
    return lm

@guidance
def add_program_and_tests(lm, name, program, args, expected):
    """Helper to format program and tests nicely"""
    lm += f'Here is an implementation of {name}:\n'
    lm += '```python\n'
    lm += program + '\n'
    lm += reconstruct_tests(name, args, expected) + '```\n'
    return lm

def format_program_and_tests(name, program, args, expected):
    tests = ''.join(f'   assert {name}({arg}) == {e}\n' for arg, e in zip(args, expected))
    return f'Here is an implementation of {name}:\n```python\n{program}\ndef test_{name}():\n{tests}```\n'

def baseline_and_tests(lm, prompt, program, args, expected):
    r = re.findall(r'def (.*?)\(', prompt)
    name = r[-1]
    lm2 = lm.set('args', args)
    lm2 = lm2.set('expected', expected)
    lm2 = lm2.set('program', program)
    lm2 = lm2.set('name', name)
    lm2 += format_program_and_tests(name, program, args, expected)
    return lm2

lm = baseline_and_tests(model, prompt, baseline_program, args, expected)

lm['args'], lm['expected']

(['3, 4, 5', '1, 2, 10'], ['6.00', '-1'])

Now, if we have a generated program and a set of tests, we can write a guidance function that runs the tests and outputs the results:

# Helper function to load the program
def load_program(name, program):
    error = None
    try:
        exec(program, globals())
        fn = eval(name)
    except Exception as e:
        fn = None
        error = e
    return fn, error

# Tolerance when x and y are floats
def equals(x, y):
    if isinstance(x, float) and isinstance(y, float):
        return abs(x - y) < 0.00001
    else:
        return x == y

@guidance
def run_tests(lm, name, program, args, expected):
    fn, error = load_program(name, program)
    all_pass = True
    lm += 'Running the test(s) above gives:\n'
    for arg, e in zip(args, expected):
        # Reconstruct the test
        lm += f'assert {name}({arg}) == {e}\n'
        try:
            arg = eval(arg)
            expected_result = eval(e)
        except:
            continue
        try:
            if isinstance(arg, tuple):
                r = fn(*arg)
            else:
                r = fn(arg)
        except Exception as ex:
            r = ex
        if equals(r, expected_result):
            lm += 'Assertion passed.\n'
        else:
            all_pass = False
            lm += f'Assertion failed.\n'
            lm += f'Expected: {e}\n'
            lm += f'Actual: {r}\n'
        lm += '---\n'
    lm = lm.set('all_pass', all_pass)
    return lm

model + run_tests(lm['name'], lm['program'], lm['args'], lm['expected'])

<guidance.models._transformers.Transformers at 0x12af1d760>

Now, we can put this all together into a function that gets the LM to rewrite the program when the tests don’t work:

def run_tests_and_fix(lm, prompt, program, args, expected):
    lm2 = baseline_and_tests(lm, prompt, program, args, expected)
    name, program, args, expected = lm2['name'], lm2['program'], lm2['args'], lm2['expected']
    i = 0
    # Try this at most 3 times
    while i != 3:
        i += 1
        lm2 += run_tests(name, program, args, expected)
        # Passing the tests, I can stop.
        if lm2['all_pass']:
            break
        lm2 += f'\n'
        # Get the model to think about what's wrong
        lm2 += f'My implementation of {name} is wrong, because''' + gen(stop='\n') + '\n'
        lm2 += f'In order to fix it, I need to''' + gen(stop='\n') + '\n'
        lm2 += f'Here is a fixed implementation:\n'
        # Write a new program
        lm2 += '```python\n' + prompt + gen(max_tokens=200, stop_regex=r'```|if __name__|def test|\n[^\s]', name='program')
        lm2 += '```\n'
        # Reset the slate, start over with new program
        program = prompt + lm2['program']
        lm2 = lm + format_program_and_tests(name, program, args, expected)
        lm2 = lm2.set('program', program)
        lm + 'ae' + gen(max_tokens=10)
    return lm2

lm = run_tests_and_fix(model, prompt, baseline_program, args, expected)

program = lm['program']
exec(program)
print(triangle_area(3, 4, 5))
print(triangle_area(1, 2, 10))

6.0
-1

In this particular case, having more rounds allows the model to fix its program on the unit tests. Does it also result in a program that passes the evaluation tests?

eval_program(program, idx)

True

Yes indeed.

In conclusion, this notebook illustrates how easy it is to guide generation depending on what previous generations are (e.g. the test results depend on the current version of the code.)