add json grammar constraint

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 
 from transformers import GPT2Tokenizer, AutoModelForCausalLM
 import numpy as np
+from transformers_cfg.grammar_utils import IncrementalGrammarConstraint
+from transformers_cfg.generation.logits_process import GrammarConstrainedLogitsProcessor
 
 MODEL_NAME = "gpt2"
 
@@ -13,6 +15,12 @@ if __name__ == "__main__":
         tokenizer.pad_token_id = tokenizer.eos_token_id
         model.config.pad_token_id = model.config.eos_token_id
 
+    # Load json grammar
+    with open("json_minimal.ebnf", "r") as file:
+        grammar_str = file.read()
+    grammar = IncrementalGrammarConstraint(grammar_str, "root", tokenizer)
+    grammar_processor = GrammarConstrainedLogitsProcessor(grammar)
+
     # Define your color-coding labels; if prob > x, then label = y; Sorted in descending probability order!
     probs_to_label = [
         (0.1, "p >= 10%"),
@@ -33,7 +41,7 @@ if __name__ == "__main__":
         """
         inputs = tokenizer([prompt], return_tensors="pt")
         outputs = model.generate(
-            **inputs, max_new_tokens=50, return_dict_in_generate=True, output_scores=True, do_sample=True
+            **inputs, max_new_tokens=20, return_dict_in_generate=True, output_scores=True, logits_processor=[grammar_processor]
         )
         # Important: don't forget to set `normalize_logits=True` to obtain normalized probabilities (i.e. sum(p) = 1)
         transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)
@@ -72,7 +80,7 @@ if __name__ == "__main__":
 
         with gr.Row():
             with gr.Column():
-                prompt = gr.Textbox(label="Prompt", lines=3, value="Today is")
+                prompt = gr.Textbox(label="Prompt", lines=3, value="This is a valid json string for http request:")
                 button = gr.Button(f"Generate with {MODEL_NAME}, using sampling!")
             with gr.Column():
                 highlighted_text = gr.HighlightedText(

json_minimal.ebnf

@@ -0,0 +1,16 @@
+
+
+root   ::= object
+
+object ::= " {" ws ( string ":" ws value ("," ws string ":" ws value)* )? ws "}"
+
+value  ::= object | array | string | number | ("true" | "false" | "null") ws
+
+array  ::= "[" ws ( value ("," ws value)* )? "]" ws
+
+string ::= "\"" [a-zA-Z0-9]* "\"" ws
+
+number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws
+
+
+ws ::= ([ \t\n] ws)?

requirements.txt

@@ -1,2 +1,3 @@
 torch
 transformers>=4.26
+transformers-cfg==0.2.0