small optim

generate.py

@@ -107,12 +107,13 @@ def stream_json_objects_from_batched_tokens_generator(batched_tokens_generator:
     for tokens_batch in chain([first_batch], batched_tokens_generator):
         for stream_idx, token in enumerate(tokens_batch):
             streams[stream_idx] += token
-            try:
-                for stream_sample in islice(ijson.items(StringIteratorIO(streams[stream_idx].__iter__()), json_field + ".item", buf_size=1), skips[stream_idx], None):
-                    yield stream_sample
-                    skips[stream_idx] = +1
-            except ijson.IncompleteJSONError:
-                pass
+            if '"' in token or "}" in token:
+                try:
+                    for stream_sample in islice(ijson.items(StringIteratorIO(streams[stream_idx].__iter__()), json_field + ".item", buf_size=1), skips[stream_idx], None):
+                        yield stream_sample
+                        skips[stream_idx] = +1
+                except ijson.IncompleteJSONError:
+                    pass
 
 
 def stream_jsonl_file(filename: str, prompt: str, columns: list[str], seed: int, size: int) -> Iterator[str]: