Repository: guidance-ai/llguidance (Rust)

Issue:
Inconsistent handling of the EOS token
When I give a list of tokens, ending with EOS token, to try_consume_tokens, then it accepts all of them. But when I give the tokens one after another, then the EOS token is rejected. Is this an expected behaviour? I would rather expect that no matter how I split a list of tokens, the behaviour should be the same.

```
import llguidance, llguidance.hf, transformers
grammar_str = 'root ::= "a"'
model_id = "hsultanbey/codegen350multi_finetuned"
tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
ll_grammar = llguidance.grammar_from("grammar", grammar_str)
ll_tokenizer = llguidance.hf.from_tokenizer(tokenizer)
ll_matcher = llguidance.LLMatcher(ll_tokenizer, ll_grammar)
a = tokenizer("a")['input_ids'][0]
eos = tokenizer.eos_token_id

print(ll_matcher.try_consume_tokens([a, eos]))  # returns 2
ll_matcher.reset()
print(ll_matcher.try_consume_tokens([a]))  # returns 1
print(ll_matcher.try_consume_tokens([eos]))  # returns 0, but I'd expect 1
```


Discussion:
Thanks for pointing this out. As you can see, I've got a repro case in Rust as well.... I'm just trying to figure out exactly why this is occurring (and whether the behaviour is intended).
Stepping through some code, I have a better diagnosis of _what_ is happening, but I'm not closer to _why_.

First, useful links:
- [`try_consume_tokens()`](https://github.com/guidance-ai/llguidance/blob/94fa39128ef184ffeda33845f6d333f332a34b4d/parser/src/matcher.rs#L167) (the `with_inner()` part is just about ensuring that any panics are trapped and reported back as regular errors)
- [`check_stop()`](https://github.com/guidance-ai/llguidance/blob/94fa39128ef184ffeda33845f6d333f332a34b4d/parser/src/tokenparser.rs#L866)
- [`validate_token()`](https://github.com/guidance-ai/llguidance/blob/94fa39128ef184ffeda33845f6d333f332a34b4d/parser/src/tokenparser.rs#L355)

Inside `try_consume_tokens()`, the supplied tokens are iterated over, with `validate_token()` being called on each first. From the link, you can see that the first thing that the latter does is check to see if the parser is in a 'stopped' state. However, this state is only updated by `check_stop()`, which occurs after all the token iteration in `try_consume_tokens()` is complete. So when consuming "a&lt;eos&gt;" the parser doesn't mark itself as stopped until both have been consumed. However, when "a" is consumed by itself, the parser still gets marked as stopped (since the grammar is satisfied and no further tokens are permissible), which means that when the attempt to consume "&lt;eos&gt;" is made, it is immediately bounced by `validate_token()`.

I don't know if this behaviour is intended. The EOS token is obviously special and could be implicitly added to every grammar. It is possible (I do not know enough about parsers in general or this one in particular) that there are a number of desirable-but-mutually-contradictory properties, and this has been chosen as the resolution of least surprise.

Side note: with a grammar of `start: "a"*` then the consumption of tokens works as expected.

@hudson-ai can you add anything, based on your knowledge on the parser internals?
Some observations:

The `infoln` in `check_stop` says `only eos token allowed, stopping` - so it suggest that even after stopping the eos token should be still allowed.

Somehow relevant comment in `parser/src/tokenparser.rs`, line 938, about exiting when the parser is stopped: `This is optional; if you don't check, compute_mask() will return an error when it cannot continue anymore`.

In `parser/src/earley/parser.rs`, function `validate_tokens`, we can see that the function explicitly handles the eos token (actually, it has to, because the eos token can come also when the parser is not stopped). It works well also in a situation when the parser is (should be) stopped, as we could see while calling `try_consume_tokens()` with a longer sequence of tokens.

My proposal of fixes - two alternatives (with 1. making more sense to me):

1) The way how the "stopped" is used internally suggests that a parser should be stopped only if it cannot accept any further tokens. So the parser should switch to stopped only after eos or error, but not if `is_accepting && !can_advance` (it can still accept eos then). Such a change would solve this issue, and would also make the result of `is_stopped()` more meaningful (would correct its semantics). Note that calling `compute_mask()` or similar functions when `!is_stopped()` would still work well (allowed_tokens is non-zero for the eos token).

2) But if you do not want to change the result of `is_stopped()`, another solution is to allow accepting the eos token after being stopped, if the stop reason is `StopReason::NoExtension`. However, this is checked three times (file `parser/src/tokenparser.rs`): in `validate_token` (line 355); then inside `check_initialized` (line 345); and then inside `validate_tokens_raw` (line 421). You can alter these checks to allow execution of interiors of the function in this case (stop reason is `StopReason::NoExtension`).

Hey @pparys  thanks for continuing to push on this! I just opened a [PR](https://github.com/guidance-ai/llguidance/pull/244) that ensures that the behavior is consistent between the two branches, but I adopted a different solution than either of your proposals here.

Essentially, I made `try_consume_tokens` always behave as if you passed the tokens one-at-a-time. I did NOT change the semantics of `is_stopped` for this case (i.e. an EOS token will not be allowed after passing the `a` token in your example above`).

I'm actually open to your suggestion either allowing the EOS token after stopping, but I considered that outside of the scope of the PR (considering the one-at-a-time behavior to be more "canonical"). @mmoskal any thoughts on this?
@hudson-ai LGTM!



Patch:
diff --git a/parser/src/matcher.rs b/parser/src/matcher.rs
index e1b09cae..9da4ee1b 100644
--- a/parser/src/matcher.rs
+++ b/parser/src/matcher.rs
@@ -166,13 +166,12 @@ impl Matcher {
         self.with_inner(|inner| {
             for (idx, &t) in tokens.iter().enumerate() {
                 if !inner.parser.validate_token(t)? {
-                    let _ = inner.parser.check_stop()?;
                     return Ok(idx);
                 }
                 let bt = inner.parser.consume_token(t)?;
+                let _ = inner.parser.check_stop()?;
                 ensure!(bt == 0, "unexpected backtracking");
             }
-            let _ = inner.parser.check_stop()?;
             Ok(tokens.len())
         })
     }


Test patch:
diff --git a/sample_parser/tests/test_raw_parser.rs b/sample_parser/tests/test_raw_parser.rs
index de9a1734..f67dceff 100644
--- a/sample_parser/tests/test_raw_parser.rs
+++ b/sample_parser/tests/test_raw_parser.rs
@@ -329,3 +329,29 @@ fn test_try_consume_after_stop() {
     }
     unreachable!();
 }
+
+#[test]
+/// Test that try_consume_tokens accepts a consistent number of tokens whether
+/// 1. the EOS token is included in the input
+/// 2. the EOS token is not included, but is consumed separately afterwards
+///
+/// Note that this test is not opinionated about whether the EOS token is accepted
+/// in either case, just that the total number of tokens consumed is the same.
+fn test_try_consume_eos_consistency() {
+    let lark = r#"start: "a""#;
+    let parser = make_parser(lark);
+    let tokens = get_tok_env().tokenize("a");
+    let eos = get_tok_env().eos_token();
+    let tokens_with_eos = [tokens.as_slice(), &[eos]].concat();
+
+    let mut matcher = Matcher::new(Ok(parser));
+    let n_consumed_all = matcher.try_consume_tokens(&tokens_with_eos).unwrap();
+
+    matcher.reset().unwrap();
+
+    let n_consumed_no_eos = matcher.try_consume_tokens(&tokens).unwrap();
+    assert!(n_consumed_no_eos <= n_consumed_all);
+    let eos_consumed = matcher.try_consume_tokens(&[eos]).unwrap();
+    assert!(eos_consumed <= 1);
+    assert_eq!(n_consumed_no_eos + eos_consumed, n_consumed_all);
+}