[Request] Constrained Generation

[scottwey]

Given the current structure of candle-vllm, how difficult would it be to add constrained generation, similar to lm-format-enforcer?

I'm happy to help here however I can.

[EricLBuehler]

Hi @scottwey, thank you for bringing that up. Implementing methods which would affect sampling such as Constrained Generation should be doable. All one would need to do is inject code at the correct location to implement the Constrained Generation: perhaps an elegant Fn pointer would be best. This is the spot that I would direct you to:

candle-vllm/src/openai/pipelines/llm_engine.rs

Lines 129 to 146 in e6c9fe4

	let logits = self.pipeline.forward(
	tokens,
	positions,
	Some(&*self.cache_engine.get_kv_cache()),
	metadata,
	)?;
	let result = self.pipeline.sample(logits, &sampling_params, &seqs)?;
	for (result, (_, seq)) in zip(result, seqs) {
	match result {
	Either::Left(logprobs) => {
	seq.deref_mut().add_token(logprobs);
	}
	Either::Right(finish_reason) => {
	seq.deref_mut().set_finish_reason(finish_reason)
	}
	}
	}

As you can see, I sample the logits and then add the result to the SequenceGroup. I am not familiar with the implementation of Constrained Generation, but after reading their README, I could imagine that you would add the implementation in this region.

Please let me know if you would be interested in implementing this.

[scottwey]

@EricLBuehler I will try to take a crack at this as soon as I get some time. Thank you for the guidance. :)

[EricLBuehler]

@scottwey, I am currently working on mistral.rs. It has a simpler sampling API and overall file structure, so perhaps you could take a look there? Feel free to raise an issue for further guidance.

[EricLBuehler]

Please see EricLBuehler/mistral.rs#59 where we are developing model grammar support. If you have any questions, please feel free to reopen!