Fix test_save_stop

Author: slundberg

guidance/models/_model.py

@@ -957,7 +957,10 @@ def __call__(self, grammar, max_tokens=1000000, n=1, top_p=1, temperature=0.0, e
                 #     self._cache_state["new_token_ids"].append(sampled_token_ind)
 
                 # capture the named groups from the parse tree
-                captured_data, captured_log_prob_data = parser.get_captures()
+                new_captured_data, new_captured_log_prob_data = parser.get_captures()
+                captured_data.update(new_captured_data)
+                captured_log_prob_data.update(new_captured_log_prob_data)
+
                 # we have no valid log prob data if we didn't compute it
                 yield new_bytes[hidden_count:], is_generated, new_bytes_prob, captured_data, captured_log_prob_data, token_count - last_token_count
                 last_token_count = token_count
@@ -969,7 +972,9 @@ def __call__(self, grammar, max_tokens=1000000, n=1, top_p=1, temperature=0.0, e
                 out = new_bytes[hidden_count:]
                 if len(out) > 0:
                     # capture the named groups from the (partial) parse tree, 
-                    captured_data, captured_log_prob_data = parser.get_captures()
+                    new_captured_data, new_captured_log_prob_data = parser.get_captures()
+                    captured_data.update(new_captured_data)
+                    captured_log_prob_data.update(new_captured_log_prob_data)
                     yield out, is_generated, new_bytes_prob, captured_data, captured_log_prob_data, token_count - last_token_count # note that we don't capture groups until a complete parse right now...
                     last_token_count = token_count
                     hidden_count = 0