Extract generation logits

[quang-anh-nguyen]

Hello I am new to outlines and trying to use it for generation with regex. I would like to generate a response containing the class name for a text classification task. The class names are included in the list CLASSES

valid_charaters = r"[A-Za-z0-9 " + string.punctuation + r"]*"
regex = valid_charaters + r"(?:" + "|".join(CLASSES).lower() + r")"  + valid_charaters
generator = outlines.generate.regex(model, regex)

How ever, I want to inspact the logit produced by the FSM at every generation step (or state), after get_allowed_tokens and bias_logits are run at line 81 in outlines/outlines/generate/generator.py.

Is there a simple solution to this demand or I need to change the forked code ?

Thank you.