Saibo Geng commited on
Commit
805081b
1 Parent(s): 7a24b1b

instantiate a new grammar object for each call

Browse files
Files changed (1) hide show
  1. app.py +8 -7
app.py CHANGED
@@ -15,12 +15,6 @@ if __name__ == "__main__":
15
  tokenizer.pad_token_id = tokenizer.eos_token_id
16
  model.config.pad_token_id = model.config.eos_token_id
17
 
18
- # Load json grammar
19
- with open("json_minimal.ebnf", "r") as file:
20
- grammar_str = file.read()
21
- grammar = IncrementalGrammarConstraint(grammar_str, "root", tokenizer)
22
- grammar_processor = GrammarConstrainedLogitsProcessor(grammar)
23
-
24
  # Define your color-coding labels; if prob > x, then label = y; Sorted in descending probability order!
25
  probs_to_label = [
26
  (0.1, "p >= 10%"),
@@ -40,8 +34,15 @@ if __name__ == "__main__":
40
  Given the prompt (text), return a list of tuples (decoded_token, label)
41
  """
42
  inputs = tokenizer([prompt], return_tensors="pt")
 
 
 
 
 
 
 
43
  outputs = model.generate(
44
- **inputs, max_new_tokens=20, return_dict_in_generate=True, output_scores=True, logits_processor=[grammar_processor]
45
  )
46
  # Important: don't forget to set `normalize_logits=True` to obtain normalized probabilities (i.e. sum(p) = 1)
47
  transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)
 
15
  tokenizer.pad_token_id = tokenizer.eos_token_id
16
  model.config.pad_token_id = model.config.eos_token_id
17
 
 
 
 
 
 
 
18
  # Define your color-coding labels; if prob > x, then label = y; Sorted in descending probability order!
19
  probs_to_label = [
20
  (0.1, "p >= 10%"),
 
34
  Given the prompt (text), return a list of tuples (decoded_token, label)
35
  """
36
  inputs = tokenizer([prompt], return_tensors="pt")
37
+
38
+ # Load json grammar and create a GrammarConstrainedLogitsProcessor for each call
39
+ with open("json_minimal.ebnf", "r") as file:
40
+ grammar_str = file.read()
41
+ grammar = IncrementalGrammarConstraint(grammar_str, "root", tokenizer)
42
+ grammar_processor = GrammarConstrainedLogitsProcessor(grammar)
43
+
44
  outputs = model.generate(
45
+ **inputs, max_new_tokens=50, repetition_penalty=1.1, return_dict_in_generate=True, output_scores=True, logits_processor=[grammar_processor]
46
  )
47
  # Important: don't forget to set `normalize_logits=True` to obtain normalized probabilities (i.e. sum(p) = 1)
48
  transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)