Organize the output types

[rlouf]

With v1 users will be able to specify the output structure of the text generated by LLMs using Python types, as if they were annotating the corresponding function. Non exhaustively:

from typing import Union, Literal
from outlines import Generator
from pydantic import BaseModel

class User(BaseModel):
    user_id: int
    first_name: str
    last_name: str

class Client(BaseModel):
    client_id: int
    first_name: str
    last_name: str

Generator(model, User)
Generator(model, list[User])
Generator(model, int)
Generator(model, list[int])
Generator(model, Literal["a", "b", "c"])
Generator(model, Union[User, Client])
Generator(model, dict[str, str])
Generator(model, tuple[int, int])

Internally we have three different ways to constrain the output type:

1. Using regex-based structured generation
2. Using grammar-based structured generation
3. Passing the type to an API model provider

So we need to internally map Python types to these different ways. I think this is where the distinction between APIGenerator and LocalGenerator comes in handy.

APIGenerator

We currently create Chocie, JsonSchema objects that we then pass to the respective model classes. I think this is an antipattern, which can be seen in the GeminiTypeAdapter class:

@format_output_type.register(Json)
def format_json_output_type(self, output_type):
    """Gemini only accepts Pydantic models and TypeDicts to define the JSON structure."""
    if issubclass(output_type.definition, BaseModel):
        return {
            "response_mime_type": "application/json",
            "response_schema": output_type.definition,
        }
    elif isinstance(output_type.definition, _TypedDictMeta):
        return {
            "response_mime_type": "application/json",
            "response_schema": output_type.definition,
        }
    else:
        raise NotImplementedError

where a much cleaner version, which does not require to introduce intermediate code, would be:

@format_output_type.register(type(BaseModel))
def format_json_output_type(self, output_type):
        return {
            "response_mime_type": "application/json",
            "response_schema": output_type,
        }

@format_output_type.register(_TypedDictMeta)
def format_json_output_type(self, output_type):
        return {
            "response_mime_type": "application/json",
            "response_schema": output_type,
        }

I thus suggest that, for API model providers, we just forward the output type passed by the user and handle the conversion in the respective XTypeAdapter classes.

LocalGenerator

Here we need to distinguish what should be handled with regex-based and CFG-based structured generation, then build the respective logits processors. For instance:

from outlines.types.regex import to_regex
from outlines.types.cfg import to_cfg


if isinstance(output_type, Cfg):
    cfg_str = to_cfg(output_type)
    logits_processor = CFGLogitsProcessor(...)
else:
    regex_str = to_regex(output_type)
    logits_processor = RegexLogitsProcessor(...)

And then handle conversion in to_regex:

def to_regex(output_type):
    case output_type:
       match JsonSchema():
            return ...
       match list():
            return ...
       match bool():
            return ...

etc.

Note: We probably want the model instance to be responsible for building the logits processor, since different models can have different backends (NumPy, Torch, MLX, JAX, etc.).

[RobinPicard]

I think the to_regex function for the LocalGenerator is a very good idea to be able to easily handle numerous types (including native types) for local models. An issue I see with what you propose for the API models is that we would lose the very straightforward switching between local and API models. I would not be able to use the same output types for both and it may be confusing for users that the Json type proposed by outlines does not work even if the API model supports json-based constrained generation.

[rlouf]

the Json type proposed by outlines does not work even if the API model supports json-based constrained generation.

How is that so? There would be a unified user interface.

[RobinPicard]

I may have misunderstood a part of it. Would your proposition include getting rid of the custom types Json and Choice? If so, my comment above does not apply. Otherwise we would need to specify that those should only be used with local models, but that's not a big deal actually.

I thought that the big advantage of going through an intermediary Json class was to translate all the different ways of expressing that type into the most basic one that is accepted by all models that support that type. In the case of Json that would a JSON schema string (that can then be turned into a regex for local models). But as Gemini does not support JSON schema string, there's no such thing as a shared most basic type and this idea crumbles.

So, thinking more about it, I agree with your proposition and think it's what makes the most sense. The only potential challenge I can see is properly conveying information on what types can be used for local models and what they do. The list type could be an ambiguous case for some users for instance as a list of strings and a list of a single type would mean different things.

[rlouf]

I may have misunderstood a part of it. Would your proposition include getting rid of the custom types Json and Choice?

Yes, we would get read of Json and Choice. As you said below, they're leaky abstractions in this context.

The list type could be an ambiguous case for some users for instance as a list of strings and a list of a single type would mean different things.

I agree, since we used lists of strings in Choice to specify multiple choices. I think that if we explain the right mental model "think of it as the output type of a Python function" the transition shouldn't be too hard?

[rlouf]

List["a", "b"] is not the correct type in Python to express "either a or b". It's Literal["a", "b"]