batched generation

generate.py

@@ -3,6 +3,7 @@ import json
 import logging
 import regex
 import time
+from itertools import chain, islice
 from pathlib import Path
 from typing import Annotated, Iterator
 
@@ -22,14 +23,16 @@ 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 = models.transformers(model_id, device=device)
+    model_id = "Qwen/Qwen1.5-0.5B-Chat"
+    batch_size = 4
 else:
     device = "cuda"
-    model = models.transformers(model_id, device=device)
+    model_id = "google/gemma-2b-it"
+    batch_size = 20
+
+model = models.transformers(model_id, device=device)
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 sampler = PenalizedMultinomialSampler()
@@ -95,6 +98,23 @@ 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
+            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=})")
@@ -134,7 +154,8 @@ def stream_jsonl_file(filename: str, prompt: str, columns: list[str], seed: int,
         tokenize=False,
         add_generation_prompt=True
     )
-    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)):
+    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)):
         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)")

gradio_app.py

@@ -6,11 +6,11 @@ import io
 import pandas as pd
 import spaces
 
-from generate import model_id, stream_jsonl_file
+from generate import model_id, stream_jsonl_file, batch_size
 
-MAX_SIZE = 20
+MAX_SIZE = 20 * batch_size
 DEFAULT_SEED = 42
-DEFAULT_SIZE = 3
+DEFAULT_SIZE = 5 * batch_size
 
 @spaces.GPU(duration=120)
 def stream_output(query: str, continue_content: str = ""):