Add grammar to transformers and _HF loaders (#4091)

modules/grammar.py

@@ -0,0 +1,33 @@
+from torch_grammar import GrammarSampler
+from transformers.generation.logits_process import LogitsProcessor
+
+from modules import shared
+
+sampler = None
+grammar = None
+grammar_string = ''
+
+
+class GrammarLogitsProcessor(LogitsProcessor):
+    def __init__(self, string):
+
+        global sampler, grammar, grammar_string
+
+        if string != grammar_string:
+            grammar_string = string
+            if string.strip() != '':
+                string = string.strip() + '\n'
+                sampler = GrammarSampler(string, 'root', shared.tokenizer)
+            else:
+                sampler = None
+
+        if sampler is not None:
+            grammar = sampler.logits_processor()
+        else:
+            grammar = None
+
+    def __call__(self, input_ids, scores):
+        if grammar is not None:
+            scores = grammar(input_ids, scores)
+
+        return scores