Filtering tokens with logit preprocessors

You may want to employ a specific pattern when generating text using your language model. One way to achieve this is by using logit preprocessors, which are only supported in return nodes of chains. You can pass a list of logit preprocessors as an argument to the return node constructor. In the following example, we will use a logit preprocessor to ensure that only tokens containing numbers are outputted by the model.

First, we define a set number_tokens within our new logit preprocessor class. The set contains all the tokens that contain only digits. We loop through each token ID and its corresponding token string using session.get_vocab().items(), which returns a dictionary with keys as token IDs and values as corresponding token strings. If the regular expression r"^ *[0-9]+$" matches the token string, it is added to the set number_tokens.

from langdash.logit_preprocessors import LogitPreprocessor

class DigitFilter(LogitPreprocessor):
  def __init__(self, session):
    self.number_tokens = set()
    for tokstr, tokid in session.get_vocab().items():
      if re.match(r"^ *[0-9]+$", tokstr) is not None:
        self.number_tokens.add(tokid)

We can then define our logit preprocessor callback, which takes in a list of input ids, and a logit vector for the current token. The function loops through each logit and checks if its corresponding token ID is not in the set number_tokens. If it is, the logit is set to -inf. This is done to mask any other tokens that isn’t a digit.

class DigitFilter(LogitPreprocessor):  
  def __call__(self, input_ids, logits):
    for token in range(len(logits)):
      if token not in self.number_tokens:
        logits[token] = -inf

Finally, we can build a chain that guides our language model to generate a sequence of numbers ranging from 1 to 10, and output the result.

chain = ld.chain(
  returns={"numbers": str},
  nodes=[
    "Here's a list of numbers from 1 to 10: ",
    ld.returns(
      "numbers",
      "",
      logit_preprocessors=[DigitFilter],
      infer_args=InferArgs(
        min_new_tokens=10,
        max_new_tokens=10,
      )
    )
  ]
)

for response in chain.stream(session):
  print(response)