Vectorize function calls

[rlouf]

In this PR I introduce a outlines.vectorize function that works similarly to numpy.vectorize, and allows to turn functions into functions that take arrays as an input. The major differences with numpy.vectorize are:

1. The possibility to batch the execution. Launching too many concurrent calls at the same time, doing inference locally with a batch size too large may not be acceptable and thus we need to batch the parallel calls.
2. It will execute coroutines concurrently for better performance.

This function will be used internally for model calls, so we can pass arrays of arbitrary shapes as inputs and get an array of answers. Usual broadcasting rules apply:

import outlines.models as models

complete = models.text_completion.openai("text-davinci-003")

prompts = np.array([["Hi", "Hello"], ["Bye", "Ciao"]])
answer = complete(prompts)
print(answer.shape)
# (2, 2)

answer_samples = complete(prompts, num_samples=10)
print(answer_samples.shape)
# (2, 2, 10)

I made the choice here to execute the coroutines in a new event loop before returning. This allows to keep the synchronous user-facing API for now. Of course, this restricts the potential speedups we get by using asyncio and we will eventually need to move towards an explicitly async API. But that would be too big of a change for a single PR.

• Vectorize scalar functions and coroutines (OpenAI API, simple tools)
• Vectorize functions with an arbitrary signature (HF models)
• Vectorize coroutines with an arbitrary signature (OpenAI with several samples)
• Execute coroutines in a new event loop before returning
• Replace the OpenAI calls with async ones and use vectorize
• Use vectorize for HF model calls

[rlouf]

Update: Vectorization is currently implemented, although it needs more documentation and tests to make the PR mergeable.

[brandonwillard]

Looks like we could use some tests for OpenAI, but I imagine that's not particularly straightforward.

Aside from that, just some questions/comments; otherwise, looks good to me.

[brandonwillard]

Do we want to provide an interface that allows the caller to manage the event loop (e.g. for use in external async applications)?

[rlouf]

I think the function that returns the new sequence should be async, but I didn't want to break the API just yet.

[brandonwillard]

Same consideration here.

[rlouf]

Looks like we could use some tests for OpenAI, but I imagine that's not particularly straightforward.

We need to refactor the code to make it easier to test the code that relies on the OpenAI API. That's the thing that annoys me the most currently.

[rlouf]

We might want a slightly different behavior for HuggingFace models: in this case we're better off flattening the input arrays so inference can be run in batch. This does not make this PR any less useful, but probably means that we'll need a different mechanism for these models. Will open an issue once this is merged.