Feature: support structured generation through Outlines

[RobinPicard]

Adding support to structured generation to LangExtract would make a lot of sense to make sure the model's output conforms to the desired format of the user. For that purpose, I propose integrating with the Outlines library that specializes in structured generation.

I think there are several reasons supporting the idea of relying on Outlines instead of implementing it ourselves in LangExtract:

• Managing the output formats specific to each model is cumbersome to implement and maintain for LangExtract.
• The diversity in expected formats and keyword arguments means that users cannot easily switch from one model to another while Outlines proposes a unified interface for structured generation.
• Local models (transformers, llama_cpp, mlx...) do not have an easy to use argument for structured generation and instead rely on providing a logits processor. That means that LangExtract would have to handle the creation of the logits processor. Instead, Outlines would take care of that and the user could just provide the same output format they would give to an API-based model.
• Outlines exposes a high-level Python interface that is convenient to work with.

Let me know if you think it's something worth looking into. I would be glad to open a PR for it.

[aksg87]

Hi @RobinPicard, I would be very interested in exploring this, maybe starting with the new plugin ecosystem introduced in #97. In particular, I would also be very interested in showing a demo of this with gemma2:2b on some long context extractions.

[the-vampiire]

i would like to try this. @aksg87 any quick advice on how you decouple the plugin from the schema? as i see it the plugin can handle the lower level bits (using transformers + logit control) but that should remain generic (not schema specific).

where then would a pydantic schema fit in to connect it with the OutlinesProvider?

[torchss]

@RobinPicard in addition to Outlines for providing structured generation instead of hoping the backend LLM does a good job of it, do you think it would be useful to add LiteLLM core provider support to langextract such that OpenAI-compatible API servers like vLLM or SGLANG support could be used?

[aksg87]

Hi @the-vampiire,

The new plugin generator script shows how to handle schema and model providers:

python scripts/create_provider_plugin.py OutlinesProvider --with-schema

The strategy LX uses is to use examples to guide schema creation, which can then be passed to the model API for better constraints during generation. The Gemini schema implementation provides a good example of this pattern.

See PR #144 for details and issue #99 for further plugin discussions.

[the-vampiire]

man this was a lot more confusing than i expected it to be. sorry for the delay, went down a lot of dead ends. i think fundamentally trying to fit a structured generation tool (more of a middleware than a provider itself) into a provider concept led to that confusion.

heres the repo if you guys want to have a look. feel free to fork or i can add you to it to make changes. i am new to both langextract and outlines so forgive the ignorance, at least its something to start off of.

i have left many ASK: and TODO: comments you can look at.

overall i feel i must have done something wrong because it cant even process the simple RJ example using phi-3 both in mini and medium.

https://github.com/the-vampiire/langextract-outlines-plugin

[RobinPicard]

Awesome, thanks a lot @the-vampiire! I think it may be easier to use the language_model_params argument to pass down all arguments that would be used in Outlines (the args to initialize the models, the output_type and the inference kwargs). That way, we would be able to benefit from the output type system of Outlines (so we would not use the examples to output type system of Langextract) and we would give full flexibility to the user regarding the creation of the Outlines model.

Here's the interface I had in mind for using Outlines with a Transformers model for instance:

result = lx.extract(
    text="Your document here",
    model_id="outlines:transformers",
    prompt_description="Extract entities",
    examples=[...],
    language_model_params={
        "model_init": [hf_model, hf_tokenizer],
        "output_type": MyPydanticClass,
        "top_k": 20,
        "repetition_penalty": 0.5,
        ...
    }
)

What do you think about it?

[the-vampiire]

i had to dig through LE source to figure this out but it looks like the simplest way to truly override and ignore examples is through the extract().model param.

i pushed up on my repo if you want to take a look, very similar to what you are describing here.

model_id = "microsoft/Phi-3-medium-4k-instruct"
tf_tokenizer = transformers.AutoTokenizer.from_pretrained(
    model_id,
    # padding_side="left",
    # trust_remote_code=True,
)
tf_model = transformers.AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="mps",
    # trust_remote_code=True,
    # pad_token_id=tf_tokenizer.eos_token_id
)

result = lx.extract(
    # text_or_documents="https://www.gutenberg.org/files/1513/1513-0.txt",
    text_or_documents="Lady Juliet gazed longingly at the stars, her heart aching for Romeo",
    prompt_description=prompt,
    examples=examples,
    max_workers=1,
    extraction_passes=3,
    # model_id="outlines",
    # extraction_passes=3,  # Multiple passes for improved recall
    # max_workers=20,  # Parallel processing for speed
    max_char_buffer=1000,  # Smaller contexts for better accuracy
    # TODO: use model param to encapsulate everything?
    model=OutlinesStructuredGenerationModel(
        model=outlines.from_transformers(tf_model, tf_tokenizer),
        output_type=SimpleExtraction,
    ),

overall id say my biggest time sink and confusion is around modeling the output type extraction schemas. theres a dissonance between langextract examples suggesting multiple output types vs being able to only provide one to outlines (along with langextract having a dependency on a { EXTRACTIONS_KEY: [OutputType] } shape downstream). would be great to get @RobinPicard @aksg87 insight into how this could be adopted.

unfortunately i still cannot get it working. no matter what i do i end up with malformed JSON that breaks with something like '{"extractions": [{"extraction_text": "Juliet", "extraction_class": " (variations on this depending on the extraction schema model provided).

this is really unexpected because using the same model (phi3 medium 4k) i have no issue processing it directly through langextract without outlines.

i ran into numerous issues trying to get this sorted out:

1. having to set os.environ["TOKENIZERS_PARALLELISM"] = "false" due to this transformers warnning showing up repeatedly

huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
        - Avoid using `tokenizers` before the fork if possible
        - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)

2. unclear how to model the output type extraction schema as a pydantic model. i tried several different ways (commented out) but even something as basic as this one still fails

class SimpleExtraction(BaseModel):
    extraction_text: str
    extraction_class: Literal["character"] | Literal["emotion"] | Literal["relationship"]

3. langextract requires a specific shape of { EXTRACTION_KEY: List[ExtractionSchema] } downstream so i had to find a way to create a dynamic wrapper using this constant key and the provided output type extraction schema class. i get a type error but it does seem to be working when i dump the derived model while debugging
4. kwargs are difficult to pass down because each outlines provider has their own set of expected kwargs. unclear if the others are as strict but transformers threw an exception when it received unexpected kwargs. for now i am just manually setting the kwargs in the generator until this is resolved

[the-vampiire]

also if you want i can fork langextract and open a (draft) PR with this work in it rather than its own repo.

it doesnt have to be merged if you guys decide its best as an external lib, but it would certainly simpl/unify discussion instead of jumping around.

[RobinPicard]

The provider plugin for Outlines has been created and is available on pypi!

Is there somewhere in langextract where I could write documentation to let users know about its existence and maybe also provide an example @aksg87?

[aksg87]

Hi @RobinPicard - that's great to hear!

You can document this in LangExtract similar to this PR: #244 so it can be listed in COMMUNITY_PROVIDERS.md.