The JSON Schema Logits Processor is a Python tool that uses PyTorch and the Transformers library to process logits based on a given JSON Schema. The tool is designed to process sequences of tokens and determine which tokens would be valid next according to a given schema.
Uses JsonSchemaLogitsProcessor
, a class that inherits from LogitsProcessor in the Transformers library.
Utilizes a tokenizer from the Transformers library.
Decodes input sequences and determines the next valid tokens based on the JSON schema.
Includes a method for validating partial JSON against a given schema.
The main class in this tool is JsonSchemaLogitsProcessor. An instance of this class can be created with a JsonSchema object and a PreTrainedTokenizer object.
from json_schema_logits_processor.schema import JsonSchema
from transformers import PreTrainedTokenizer
from json_schema_logits_processor.json_schema_logits_processor import JsonSchemaLogitsProcessor
from json_schema_logits_processor.schema import parse_schema_from_string
schema = parse_schema_from_string(
'{"type": "object", "properties": {"a": {"type": "string"}}}'
)
# Assuming `schema` is a JsonSchema object and `tokenizer` is a PreTrainedTokenizer object
processor = JsonSchemaLogitsProcessor(schema, tokenizer)
Usage in transformers models:
output = model.generate(*inputs, logits_processor=[logits_processor])
The processor decodes the input sequence and determines the next valid tokens based on the JSON schema. The returned tensor contains the scores for the valid tokens and a very low score for invalid tokens.
- PyTorch
- Transformers
Please note that this project is a work in progress and the functionality may change over time.
For more information, please review the source code or contact the author. Contributions and suggestions are welcome!
Contact Evan Richards
- Github: @evanrichards
mit