How one can constrain your model to output defined formats

On this post I’ll explain and display the concept of “structured generative AI”: generative AI constrained to defined formats. By the tip of the post, you’ll understand where and when it may possibly be used and the best way to implement it whether you’re crafting a transformer model from scratch or utilizing Hugging Face’s models. Moreover, we are going to cover a vital tip for tokenization that is very relevant for structured languages.

Considered one of the numerous uses of generative AI is as a translation tool. This often involves translating between two human languages but may include computer languages or formats. For instance, your application might have to translate natural (human) language to SQL:

Natural language: “Get customer names and emails of consumers from the US”
SQL: "SELECT name, email FROM customers WHERE country = 'USA'"

Or to convert text data right into a JSON format:

Natural language: “I'm John Doe, phone number is 555–123–4567,
my friends are Anna and Sara”
JSON: {name: "John Doe",
phone_number: "555–123–5678",
friends: {
name: [["Anna", "Sara"]]}
}

Naturally, many more applications are possible, for other structured languages. The training process for such tasks involves feeding examples of natural language alongside structured formats to an encoder-decoder model. Alternatively, leveraging a pre-trained Language Model (LLM) can suffice.

While achieving 100% accuracy is unattainable, there may be one class of errors that we are able to eliminate: syntax errors. These are violations of the format of the language, like replacing commas with dots, using table names that usually are not present within the SQL schema, or omitting bracket closures, which render SQL or JSON non-executable.

The indisputable fact that we’re translating right into a structured language signifies that the list of legitimate tokens at every generation step is proscribed, and pre-determined. If we could insert this data into the generative AI process we are able to avoid a wide selection of incorrect results. That is the concept behind structured generative AI: constrain it to an inventory of legitimate tokens.

A fast reminder on how tokens are generated

Whether employing an encoder-decoder or GPT architecture, token generation operates sequentially. Each token’s selection relies on each the input and previously generated tokens, continuing until a token is generated, signifying the completion of the sequence. At each step, a classifier assigns logit values to all tokens within the vocabulary, representing the probability of every token as the subsequent selection. The following token is sampled based on those logits.

Limiting token generation

To constrain token generation, we incorporate knowledge of the output language’s structure. Illegitimate tokens have their logits set to -inf, ensuring their exclusion from selection. As an illustration, if only a comma or “FROM” is valid after “Select name,” all other token logits are set to -inf.

For those who’re using Hugging Face, this might be implemented using a “logits processor”. To make use of it it’s worthwhile to implement a category with a __call__ method, which shall be called after the logits are calculated, but before the sampling. This method receives all token logits and generated input IDs, returning modified logits for all tokens.

I’ll display the code with a simplified example. First, we initialize the model, we are going to use Bart on this case, but this will work with any model.

from transformers import BartForConditionalGeneration, BartTokenizerFast, PreTrainedTokenizer
from transformers.generation.logits_process import LogitsProcessorList, LogitsProcessor
import torch
name = 'facebook/bart-large'
tokenizer = BartTokenizerFast.from_pretrained(name, add_prefix_space=True)
pretrained_model = BartForConditionalGeneration.from_pretrained(name)

If we would like to generate a translation from the natural language to SQL, we are able to run:

to_translate = 'customers emails from the us'
words = to_translate.split()
tokenized_text = tokenizer([words], is_split_into_words=True)
out = pretrained_model.generate(
torch.tensor(tokenized_text["input_ids"]),
max_new_tokens=20,
)
print(tokenizer.convert_tokens_to_string(
tokenizer.convert_ids_to_tokens(
out[0], skip_special_tokens=True)))

Returning

'More emails from the us'

Since we didn’t fine-tune the model for text-to-SQL tasks, the output doesn’t resemble SQL. We won’t train the model on this tutorial, but we are going to guide it to generate an SQL query. We’ll achieve this by employing a function that maps each generated token to an inventory of permissible next tokens. For simplicity, we’ll focus only on the immediate preceding token, but more complicated mechanisms are easy to implement. We’ll use a dictionary defining for every token, which tokens are allowed to follow it. E.g. The query must begin with “SELECT” or “DELETE”, and after “SELECT” only “name”, “email”, or ”id” are allowed since those are the columns in our schema.

rules = {'': ['SELECT', 'DELETE'], # starting of the generation
'SELECT': ['name', 'email', 'id'],  # names of columns in our schema
'DELETE': ['name', 'email', 'id'],
'name': [',', 'FROM'],
'email': [',', 'FROM'],
'id': [',', 'FROM'],
',': ['name', 'email', 'id'],
'FROM': ['customers', 'vendors'],  # names of tables in our schema
'customers': [''],
'vendors': [''],  # end of the generation
}

Now we’d like to convert these tokens to the IDs utilized by the model. This can occur inside a category inheriting from LogitsProcessor.

def convert_token_to_id(token):
return tokenizer(token, add_special_tokens=False)['input_ids'][0]
class SQLLogitsProcessor(LogitsProcessor):
def __init__(self, tokenizer: PreTrainedTokenizer):
self.tokenizer = tokenizer
self.rules = {convert_token_to_id(k): [convert_token_to_id(v0) for v0 in v] for k,v in rules.items()}

Finally, we are going to implement the __call__ function, which is named after the logits are calculated. The function creates a brand new tensor of -infs, checks which IDs are legitimate in keeping with the foundations (the dictionary), and places their scores in the brand new tensor. The result’s a tensor that only has valid values for the valid tokens.

class SQLLogitsProcessor(LogitsProcessor):
def __init__(self, tokenizer: PreTrainedTokenizer):
self.tokenizer = tokenizer
self.rules = {convert_token_to_id(k): [convert_token_to_id(v0) for v0 in v] for k,v in rules.items()}
def __call__(self, input_ids: torch.LongTensor, scores: torch.LongTensor):
if not (input_ids == self.tokenizer.bos_token_id).any():
# we must allow the beginning token to look before we start processing
return scores
# create a brand new tensor of -inf
new_scores = torch.full((1, self.tokenizer.vocab_size), float('-inf'))
# ids of legitimate tokens
legit_ids = self.rules[int(input_ids[0, -1])]
# place their values in the brand new tensor
new_scores[:, legit_ids] = scores[0, legit_ids]
return new_scores

And that’s it! We will now run a generation with the logits-processor:

to_translate = 'customers emails from the us'
words = to_translate.split()
tokenized_text = tokenizer([words], is_split_into_words=True, return_offsets_mapping=True)
logits_processor = LogitsProcessorList([SQLLogitsProcessor(tokenizer)])
out = pretrained_model.generate(
torch.tensor(tokenized_text["input_ids"]),
max_new_tokens=20,
logits_processor=logits_processor
)
print(tokenizer.convert_tokens_to_string(
tokenizer.convert_ids_to_tokens(
out[0], skip_special_tokens=True)))

Returning

 SELECT email , email , id , email FROM customers

The final result is a bit of strange, but remember: we didn’t even train the model! We only enforced token generation based on specific rules. Notably, constraining generation doesn’t interfere with training; constraints only apply during generation post-training. Thus, when appropriately implemented, these constraints can only enhance generation accuracy.

Our simplistic implementation falls wanting covering all of the SQL syntax. An actual implementation must support more syntax, potentially considering not only the last token but several, and enable batch generation. Once these enhancements are in place, our trained model can reliably generate executable SQL queries, constrained to valid table and column names from the schema. A Similar approach can implement constraints in generating JSON, ensuring key presence and bracket closure.

Watch out of tokenization

Tokenization is usually missed but correct tokenization is crucial when using generative AI for structured output. Nonetheless, under the hood, tokenization could make an impact on the training of your model. For instance, chances are you’ll fine-tune a model to translate text right into a JSON. As a part of the fine-tuning process, you provide the model with examples of text-JSON pairs, which it tokenizes. What is going to this tokenization seem like?

When you read “[[“ as two square brackets, the tokenizer converts them right into a single ID, which shall be treated as a very distinct class from the only bracket by the token classifier. This makes your complete logic that the model must learn — more complicated (for instance, remembering what number of brackets to shut). Similarly, adding an area before words may change their tokenization and their class ID. As an illustration:

Again, this complicates the logic the model could have to learn because the weights connected to every of those IDs could have to be learned individually, for barely different cases.

For less complicated learning, ensure each concept and punctuation is consistently converted to the identical token, by adding spaces before words and characters.

Inputting spaced examples during fine-tuning simplifies the patterns the model has to learn, enhancing model accuracy. During prediction, the model will output the JSON with spaces, which you’ll then remove before parsing.

Summary

Generative AI offers a useful approach for translating right into a formatted language. By leveraging the knowledge of the output structure, we are able to constrain the generative process, eliminating a category of errors and ensuring the executability of queries and parse-ability of knowledge structures.

Moreover, these formats may use punctuation and keywords to indicate certain meanings. Ensuring that the tokenization of those keywords is consistent can dramatically reduce the complexity of the patterns that the model has to learn, thus reducing the required size of the model and its training time, while increasing its accuracy.

Structured generative AI can effectively translate natural language into any structured format. These translations enable information extraction from text or query generation, which is a robust tool for varied applications.