116 lines
3.1 KiB
Julia
116 lines
3.1 KiB
Julia
# Code Generation - MCTS for Programming Tasks
|
|
|
|
This example shows how MCTS can guide LLM code generation by exploring different implementation strategies.
|
|
|
|
```julia
|
|
using LLMMCTS
|
|
|
|
# State represents the current state of code generation
|
|
# It includes the code written so far and the problem being solved
|
|
|
|
function code_generation_transition(state::Dict, args::NamedTuple)
|
|
current_step = get(state, :step, 0)
|
|
problem = state[:problem]
|
|
code_so_far = get(state, :code, "")
|
|
|
|
if current_step == 0
|
|
# First step: Plan the approach
|
|
new_code = """
|
|
# Function to solve: $(problem)
|
|
function solve_problem(input)
|
|
"""
|
|
newstate = Dict(
|
|
:step => 1,
|
|
:code => new_code,
|
|
:thought => "Plan the approach for: $(problem)",
|
|
:reward => 2.0,
|
|
:isterminal => false
|
|
)
|
|
progressvalue = 5.0
|
|
elseif current_step == 1
|
|
# Second step: Implement main logic
|
|
new_code = code_so_far * """
|
|
# Main logic implementation
|
|
result = input * 2 # Placeholder implementation
|
|
return result
|
|
end
|
|
"""
|
|
newstate = Dict(
|
|
:step => 2,
|
|
:code => new_code,
|
|
:thought => "Implement main function logic",
|
|
:reward => 3.0,
|
|
:isterminal => false
|
|
)
|
|
progressvalue = 7.0
|
|
elseif current_step == 2
|
|
# Third step: Add tests
|
|
new_code = code_so_far * """
|
|
|
|
# Test the function
|
|
@assert solve_problem(5) == 10
|
|
@assert solve_problem(0) == 0
|
|
println("All tests passed!")
|
|
"""
|
|
newstate = Dict(
|
|
:step => 3,
|
|
:code => new_code,
|
|
:thought => "Add unit tests to verify implementation",
|
|
:reward => 5.0,
|
|
:isterminal => true # Code generation complete
|
|
)
|
|
progressvalue = 10.0
|
|
else
|
|
newstate = Dict(
|
|
:step => current_step,
|
|
:code => code_so_far,
|
|
:thought => "Code generation complete",
|
|
:reward => 10.0,
|
|
:isterminal => true
|
|
)
|
|
progressvalue = 10.0
|
|
end
|
|
|
|
return Dict(
|
|
:newNodeKey => "code_step_$current_step",
|
|
:newstate => newstate,
|
|
:progressvalue => progressvalue
|
|
)
|
|
end
|
|
|
|
# Initial state
|
|
initialstate = Dict(
|
|
:step => 0,
|
|
:problem => "Create a function that doubles its input",
|
|
:code => "",
|
|
:reward => 0,
|
|
:isterminal => false
|
|
)
|
|
|
|
# Transition arguments
|
|
transitionargs = (max_steps = 3,)
|
|
|
|
# Run MCTS
|
|
result = runMCTS(
|
|
initialstate,
|
|
code_generation_transition,
|
|
transitionargs;
|
|
maxiterations = 20,
|
|
explorationweight = 1.0,
|
|
maxSimulationDepth = 3,
|
|
horizontalSampleExpansionPhase = 3
|
|
)
|
|
|
|
# Display results
|
|
println("Code Generation Example")
|
|
println("=======================")
|
|
println()
|
|
println("Problem: ", initialstate[:problem])
|
|
println()
|
|
println("Generated code:")
|
|
println(result.bestTerminalState[:code])
|
|
println()
|
|
println("Code generation complete! ✓")
|
|
println("Root node visits: ", result.root.visits)
|
|
```
|