Adding MLXLM, VLLM classes to LogitsGenerator type

[parkervg]

Very excited about the mlx integration in v0.0.44!

Began integrating it into a project, but noticed that type hinting was throwing some errors around the LogitsGenerator type. Essentially, it looks as though this type wasn't synced up with the creation of the two new models.

This PR just adds the two models to this type in models/__init__.py:

LogitsGenerator = Union[Transformers, LlamaCpp, ExLlamaV2Model, Mamba, MLXLM, VLLM]

However, it brought up a bigger design question I'm curious about. What was the reasoning to go with this pattern (building a custom type using Union) as opposed to an abstract LogitsGenerator class, which all future model classes subclass? This way:

• You wouldn't need to worry about first creating the model, and then going to update LogitsGenerator
• We'd have stricter constraints over what a valid model is (i.e a generate() function returning a float is a no-go)
• We could type the 'loader functions' (llamacpp, mamba, mlxlm, etc.) to explicitly return a LogitsGenerator

Thanks!

[lapp0]

It makes sense that we have an abstract base class for all models, I agree.

Began integrating it into a project, but noticed that type hinting was throwing some errors around the LogitsGenerator type. Essentially, it looks as though this type wasn't synced up with the creation of the two new models.

Could you help me understand the type hinting error? Does your downstream library have a function / class which accepts any outlines model with the LogitsGenerator type hint?

MLXLM isn't a logits generator, it isn't designed to provide logits via __call__, it only allows generate() and stream(). Is generating logits token-by-token via __call__ is this necessary for your use case? Or could you help me understand your use case better?

[parkervg]

Ahh ok, I might be misinterpreting that type then. From browsing the code, it seemed to me as though being a LogitsGenerator type was a prerequisite for being used in any of the outlines.generate functions.

For my specific usecase, I was doing something like this (abbreviated semi-pseudo code).

Currently I just use regex and choice in my project, but in the future might use others.

def predict(model: LogitsGenerator, gen_type: Literal['generate', 'regex']):
    if gen_type == "generate":
        generator = outlines.generate.choice(model, ...)
    elif gen_type == "regex":
        generator = outlines.generate.regex(model, ...)
    return generator(prompt)

So essentially, I guess my main question is - do you have a type defined within the outlines codebase currently that dictates it will be a valid first argument to a generate method? Or, would this only happening after implementing the abstract base model class you referred to?

(Side question - isn't the MLXLM class generating and processing logits in the generate_step function for tasks like regex generation? Or is your definition of a LogitsGenerator type something that explicitly generates logits via __call__?)

[lapp0]

So essentially, I guess my main question is - do you have a type defined within the outlines codebase currently that dictates it will be a valid first argument to a generate method? Or, would this only happening after implementing the abstract base model class you referred to?

Nope, but I'd like them to. We should have them as an OutlinesModel, please create an issue so we can create an abstract base class and have consistent interfaces and typing.

isn't the MLXLM class generating and processing logits in the generate_step function for tasks like regex generation?

generate_step is an iterator, it doesn't return a single steps logits like __call__ does. It's a bad name IMO. It's adapted from a function with the same name and signature in mlx-lm.

Or is your definition of a LogitsGenerator type something that explicitly generates logits via call?)

Based on the name, that seems like a sane definition for LogitsGenerator. This type isn't used anywhere else in the codebase so I'm not 100% sure what it means to be a LogitsGenerator though.