mixture.jl

## GPT3MixtureTrace ##

"""
    GPT3MixtureTrace

A trace generated by a mixture distribution over language model prompts.

`Base.getindex(trace::GPT3MixtureTrace, addr)` supports the
following values for `addr`:
- `:prompt`: The prompt provided as an argument to the language model.
- `:prior_probs`: The prior probability of each prompt.
- `:post_probs`: The posterior probability of each prompt.
- `:joint_probs`: The joint probability of each prompt with the output.
- `:tokens`: Vector of output tokens, including the stop sequence.
- `:output`: The output completion generated by the language model.
- `:output_scores`: Log probabilities of the output under each prompt.
- `:score`: Total log probability of the output, marginalized over prompts.
"""
struct GPT3MixtureTrace{T <: GenerativeFunction} <: Trace
    gen_fn::T
    prompts::Vector{String}
    probs::Vector{Float64}
    multi_trace::MultiGPT3Trace{MultiGPT3GF}
    output::String
    tokens::Vector{String}
    scores::Vector{Float64}
    score::Float64
end

get_choices(trace::GPT3MixtureTrace) = GPT3ChoiceMap(trace.output)
get_args(trace::GPT3MixtureTrace) = (trace.prompts, trace.probs)
get_retval(trace::GPT3MixtureTrace) = trace.output
get_score(trace::GPT3MixtureTrace) = trace.score
get_gen_fn(trace::GPT3MixtureTrace) = trace.gen_fn

function Base.getindex(trace::GPT3MixtureTrace, addr)
    if addr == :prompts
        return trace.prompts
    elseif addr == :prior_probs
        return trace.probs
    elseif addr == :post_probs
        return exp.(trace.scores .+ log.(trace.probs) .- trace.score)
    elseif addr == :joint_probs
        return exp.(trace.scores .+ log.(trace.probs))
    elseif addr == :tokens
        return trace.tokens
    elseif addr == :output
        return trace.output
    elseif addr == :output_scores
        return trace.scores
    elseif addr == :score
        return trace.score
    else
        throw(KeyError(addr))
    end
end

function Base.:(==)(trace1::GPT3MixtureTrace, trace2::GPT3MixtureTrace)
    return (trace1.gen_fn == trace2.gen_fn &&
            trace1.prompts == trace2.prompts &&
            trace1.probs == trace2.probs &&
            trace1.multi_trace == trace2.multi_trace &&
            trace1.output == trace2.output &&
            trace1.tokens == trace2.tokens &&
            trace1.scores == trace2.scores &&
            trace1.score == trace2.score)
end

## GPT3Mixture ##

"""
    GPT3Mixture(;
        model = "davinci-002",
        temperature = 1.0,
        max_tokens = 1024,
        stop = nothing,
        batch_size = 10,
        encoding = GenGPT3.MODEL_ENCODINGS[model],
        api_key_lookup = () -> ENV["OPENAI_API_KEY"],
        organization_lookup = () -> ENV["OPENAI_ORGANIZATION"]
    )

Mixture over [`GPT3GenerativeFunction`](@ref)s with different prompts. Takes in
a vector of prompts and a vector of probabilities, and samples from the mixture
according to the probabilities. If the probability argument is ommitted, then
the mixture is uniform over the prompts.

This is equivalent to writing a `@gen` function which first samples from a 
categorical distribution over prompts, then samples a completion to the selected
prompt. However, it operates more efficiently by using [`MultiGPT3GF`](@ref)
under the hood to make batched API calls.
"""
struct GPT3Mixture <: GenerativeFunction{String,GPT3MixtureTrace}
    multi_gf::MultiGPT3GF
end

GPT3Mixture(; kwargs...) = GPT3Mixture(MultiGPT3GF(; kwargs...))

function (gen_fn::GPT3Mixture)(
    prompts::Vector{String},
    probs::Vector{Float64} = ones(length(prompts)) ./ length(prompts)
)
    chosen_idx = categorical(probs)
    outputs = gen_fn.multi_gf(prompts[chosen_idx:chosen_idx])
    return outputs[1]
end

function simulate(gen_fn::GPT3Mixture, args::Tuple)
    # Extract arguments
    prompts = args[1]
    probs = get(args, 2, ones(length(prompts)) ./ length(prompts))
    @assert length(prompts) == length(probs)
    # Sample prompt and completion
    chosen_idx = categorical(probs)
    chosen_trace = simulate(gen_fn.multi_gf, (prompts[chosen_idx:chosen_idx],))
    output = chosen_trace.outputs[1]
    # Score completion under other prompts
    outputs = MultiGPT3ChoiceMap(fill(output, length(prompts)-1))
    unchosen_prompts = prompts[1:end .!= chosen_idx]
    multi_trace, _ = generate(gen_fn.multi_gf, (unchosen_prompts,), outputs)
    # Insert chosen trace into multi trace
    insert!(multi_trace.prompts, chosen_idx, chosen_trace.prompts[1])
    insert!(multi_trace.outputs, chosen_idx, chosen_trace.outputs[1])
    insert!(multi_trace.tokens, chosen_idx, chosen_trace.tokens[1])
    insert!(multi_trace.logprobs, chosen_idx, chosen_trace.logprobs[1])
    insert!(multi_trace.scores, chosen_idx, chosen_trace.scores[1])
    multi_trace = MultiGPT3Trace(
        multi_trace.gen_fn, multi_trace.prompts, multi_trace.outputs,
        multi_trace.tokens, multi_trace.logprobs, multi_trace.scores,
        multi_trace.score + chosen_trace.score
    )
    # Compute probability of output under prompt mixture
    scores = copy(multi_trace.scores)
    score = logsumexp(scores .+ log.(probs))
    # Construct and return trace
    tokens = copy(chosen_trace.tokens[1])
    trace = GPT3MixtureTrace(gen_fn, prompts, probs, multi_trace,
                             output, tokens, scores, score)
    return trace
end

function generate(gen_fn::GPT3Mixture, args::Tuple, constraints::ChoiceMap)
    # Check whether output is constrained
    if !has_value(constraints, OUTPUT_ADDR)
        return generate(gen_fn, args, EmptyChoiceMap())
    end
    # Extract arguments
    prompts = args[1]
    probs = get(args, 2, ones(length(prompts)) ./ length(prompts))
    @assert length(prompts) == length(probs)
    # Score completion under prompts
    output = get_value(constraints, OUTPUT_ADDR)
    outputs = MultiGPT3ChoiceMap(fill(output, length(prompts)))
    multi_trace, _ = generate(gen_fn.multi_gf, (prompts,), outputs)
    # Compute probability of output under prompt mixture
    scores = copy(multi_trace.scores)
    score = logsumexp(scores .+ log.(probs))
    # Construct and return trace
    tokens = copy(multi_trace.tokens[1])
    trace = GPT3MixtureTrace(gen_fn, prompts, probs, multi_trace,
                             output, tokens, scores, score)
    return trace, score
end

generate(gen_fn::GPT3Mixture, args::Tuple, ::EmptyChoiceMap) =
    simulate(gen_fn, args), 0.0

project(trace::GPT3MixtureTrace, selection::Selection) =
    OUTPUT_ADDR in selection ? trace.score : 0.0   

project(trace::GPT3MixtureTrace, ::EmptySelection) =
    0.0

function update(trace::GPT3MixtureTrace,
                args::Tuple, argdiffs::Tuple, constraints::ChoiceMap)
    # Check whether output is constrained
    if isempty(constraints)
        return update(trace, args, argdiffs, EmptyChoiceMap())
    elseif !has_value(constraints, OUTPUT_ADDR)
        error("Did not visit all constraints")
    end
    # Update underlying MultiGPT3Trace
    prompts = args[1]
    prompt_diffs = argdiffs[1]
    output = get_value(constraints, OUTPUT_ADDR)
    outputs = MultiGPT3ChoiceMap(fill(output, length(prompts)))
    multi_trace, _, _, _ =
        update(trace.multi_trace, (prompts,), (prompt_diffs,), outputs)
    # Compute probability of output under prompt mixture
    probs = get(args, 2, ones(length(prompts)) ./ length(prompts))
    @assert length(prompts) == length(probs)
    scores = copy(multi_trace.scores)
    score = logsumexp(scores .+ log.(probs))
    # Construct new trace
    tokens = copy(multi_trace.tokens[1])
    new_trace = GPT3MixtureTrace(trace.gen_fn, prompts, probs, multi_trace,
                                 output, tokens, scores, score)
    # Return trace and update weight, etc.
    weight = new_trace.score - trace.score
    retdiff = trace.output == output ? NoChange() : UnknownChange()
    discard = get_choices(trace)
    return new_trace, weight, retdiff, discard
end
    
function update(trace::GPT3MixtureTrace,
                args::Tuple, argdiffs::Tuple, ::EmptyChoiceMap)
    # Return same trace if arguments do not change
    prompts = args[1]
    probs = get(args, 2, ones(length(prompts)) ./ length(prompts))
    @assert length(prompts) == length(probs)
    if (argdiffs == (NoChange(), NoChange()) || 
        trace.prompts == prompts && trace.probs == probs)
        return trace, 0.0, NoChange(), EmptyChoiceMap()
    end
    # Update underlying MultiGPT3Trace
    prompt_diffs = argdiffs[1]
    multi_trace, _, _, _ =
        update(trace.multi_trace, (prompts,), (prompt_diffs,), EmptyChoiceMap())
    # Compute probability of output under prompt mixture
    scores = copy(multi_trace.scores)
    score = logsumexp(scores .+ log.(probs))
    # Construct new trace
    output = trace.output
    tokens = copy(multi_trace.tokens[1])
    new_trace = GPT3MixtureTrace(trace.gen_fn, prompts, probs, multi_trace,
                                 output, tokens, scores, score)
    # Return trace and update weight, etc.
    weight = new_trace.score - trace.score
    return new_trace, weight, NoChange(), EmptyChoiceMap()
end

function regenerate(trace::GPT3MixtureTrace,
                    args::Tuple, argdiffs::Tuple, selection::Selection)
    # Check whether output is selected
    if isempty(selection)
        return regenerate(trace, args, argdiffs, EmptySelection())
    end
    # Simulate new trace and return no change in score
    new_trace = simulate(trace.gen_fn, args)
    return new_trace, 0.0, UnknownChange()
end
    
function regenerate(trace::GPT3MixtureTrace,
                    args::Tuple, argdiffs::Tuple, ::EmptySelection)
    # Return same trace if arguments do not change
    prompts = args[1]
    probs = get(args, 2, ones(length(prompts)) ./ length(prompts))
    @assert length(prompts) == length(probs)
    if (argdiffs == (NoChange(), NoChange()) || 
        trace.prompts == prompts && trace.probs == probs)
        return trace, 0.0, NoChange(), EmptyChoiceMap()
    end
    # Update underlying MultiGPT3Trace
    prompt_diffs = argdiffs[1]
    multi_trace, _, _, _ =
        update(trace.multi_trace, (prompts,), (prompt_diffs,), EmptyChoiceMap())
    # Compute probability of output under prompt mixture
    scores = copy(multi_trace.scores)
    score = logsumexp(scores .+ log.(probs))
    # Construct new trace
    output = trace.output
    tokens = copy(multi_trace.tokens[1])
    new_trace = GPT3MixtureTrace(trace.gen_fn, prompts, probs, multi_trace,
                                 output, tokens, scores, score)
    # Return trace and update weight, etc.
    weight = new_trace.score - trace.score
    return new_trace, weight, NoChange()
end