revert batching

generate.py

@@ -3,7 +3,6 @@ import json
 import logging
 import regex
 import time
-from itertools import chain, islice
 from pathlib import Path
 from typing import Annotated, Iterator
 
@@ -23,16 +22,14 @@ logger = logging.getLogger(__name__)
 
 
 logger.warning("Loading model...")
+model_id = "google/gemma-2b-it"
+# model_id = "Qwen/Qwen1.5-0.5B-Chat"
 if torch.backends.mps.is_available():
     device = "mps"
-    model_id = "Qwen/Qwen1.5-0.5B-Chat"
-    batch_size = 1  # batching generates duplicates
+    model = models.transformers(model_id, device=device)
 else:
     device = "cuda"
-    model_id = "google/gemma-2b-it"
-    batch_size = 1  # batching generates duplicates
-
-model = models.transformers(model_id, device=device)
+    model = models.transformers(model_id, device=device)
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 sampler = PenalizedMultinomialSampler()
@@ -98,24 +95,6 @@ def samples_prommpt(filename: str, prompt: str, columns: str):
     {{ prompt }}
     """
 
-
-def stream_json_objects_from_batched_tokens_generator(batched_tokens_generator: Iterator[list[str]], json_field: str) -> Iterator[dict]:
-    first_batch = next(batched_tokens_generator)
-    batch_size = len(first_batch)
-    streams = [""] * batch_size
-    skips = [0] * batch_size
-    for tokens_batch in chain([first_batch], batched_tokens_generator):
-        for stream_idx, token in enumerate(tokens_batch):
-            streams[stream_idx] += token
-            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]:
     filename = Path(filename).stem
     logger.warning(f"stream_response({filename=}, {prompt=}, {columns=})")
@@ -155,8 +134,7 @@ def stream_jsonl_file(filename: str, prompt: str, columns: list[str], seed: int,
         tokenize=False,
         add_generation_prompt=True
     )
-    batched_samples_generator_tokens = samples_generator.stream([text] * batch_size, rng=rng)
-    json_field = list(Dataset.model_fields)[0]
-    for _, sample in zip(range(size), stream_json_objects_from_batched_tokens_generator(batched_samples_generator_tokens, json_field=json_field)):
+    samples_generator_tokens = samples_generator.stream(text, rng=rng)
+    for _, sample in zip(range(size), ijson.items(StringIteratorIO(samples_generator_tokens), "data.item", buf_size=4)):
         yield json.dumps(sample, ensure_ascii=False) + "\n"
-    logger.warning(f"stream_response({filename=}, {prompt=}, {columns=}) - Generating samples... DONE (total={time.time() - _start:.02f}s)")
+    logger.warning(f"stream_response({filename=}, {prompt=}, {columns=}) - Generating samples... DONE (total={time.time() - _start:.02f}s)")

gradio_app.py

@@ -6,11 +6,11 @@ import io
 import pandas as pd
 import spaces
 
-from generate import model_id, stream_jsonl_file, batch_size
+from generate import model_id, stream_jsonl_file
 
-MAX_SIZE = 20 * batch_size
+MAX_SIZE = 20
 DEFAULT_SEED = 42
-DEFAULT_SIZE = 5 * batch_size
+DEFAULT_SIZE = 3
 
 @spaces.GPU(duration=120)
 def stream_output(query: str, continue_content: str = ""):
@@ -87,4 +87,4 @@ with gr.Blocks() as demo:
     generate_more_button.click(stream_more_output, filename_comp, outputs)
 
 
-demo.launch()
+demo.launch()