initial commit

README.md

@@ -5,7 +5,7 @@ colorFrom: yellow
 colorTo: green
 sdk: gradio
 sdk_version: 4.25.0
-app_file: app.py
+app_file: gradio_app.py
 pinned: false
 ---
 

api.py

@@ -0,0 +1,57 @@
+import logging
+from typing import Annotated
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, StringConstraints
+from outlines import generate
+
+from generate import model, sampler, stream_file
+
+logger = logging.getLogger(__name__)
+
+
+class Status(BaseModel):
+    status: Annotated[str, StringConstraints(pattern="ok")]
+
+status_generator = generate.json(model, Status, sampler=sampler)
+
+status_stream = status_generator.stream("status:")
+status = "".join(char.strip() for char in status_stream if char.strip())
+logger.warning("Model status: " + status)
+
+
+async def stream_response(filename: str, prompt: str, columns: list[str], seed: int, size: int):
+    for chunk in stream_file(
+        filename=filename,
+        prompt=prompt,
+        columns=columns,
+        seed=seed,
+        size=size,
+    ):
+        yield chunk
+
+
+async def dummy_stream():
+    yield ""
+
+
+app = FastAPI()
+
+@app.head("/{filename}.jsonl")
+@app.get("/{filename}.jsonl")
+async def read_item(request: Request, filename: str, prompt: str = "", columns: str = "", seed: int = 42, size: int = 3):
+    if request.method == 'GET':
+        columns = [field.strip() for field in columns.split(",") if field.strip()]
+        content = stream_response(
+            filename,
+            prompt=prompt,
+            columns=columns,
+            seed=seed,
+            size=size
+        )
+    else:
+        content = dummy_stream()
+    response = StreamingResponse(content, media_type="text/jsonlines")
+    response.headers["Content-Disposition"] = f"attachment; filename={filename}.jsonl"
+    return response

fsm.py

@@ -0,0 +1,92 @@
+from copy import copy
+from functools import partial
+
+from outlines.fsm.guide import RegexGuide
+from pydantic import BaseModel
+from transformers import PreTrainedTokenizerBase
+
+
+def merge_successive_transitions(states_to_token_maps: dict[int, dict[int, int]], i, j):
+    states_to_token_maps = dict(states_to_token_maps)
+    transitions_i = {(s1, states_to_token_maps[s1][i]) for s1 in states_to_token_maps if i in states_to_token_maps[s1]}
+    transitions_j = {(s1, states_to_token_maps[s1][j]) for s1 in states_to_token_maps if j in states_to_token_maps[s1]}
+    transitions_i, transitions_j = dict(transitions_i - transitions_j), dict(transitions_j - transitions_i)
+    for s1, s2 in transitions_i.items():
+        while s2 in transitions_j:
+            s2 = transitions_j[s2]
+        if s2 != transitions_i[s1]:
+            states_to_token_maps[s1] = dict(states_to_token_maps[s1])
+            states_to_token_maps[s1][i] = s2
+    return states_to_token_maps
+
+
+def replace_transitions(states_to_token_maps: dict[int, dict[int, int]], i, j):
+    states_to_token_maps = dict(states_to_token_maps)
+    transitions_i = {(s1, states_to_token_maps[s1][i]) for s1 in states_to_token_maps if i in states_to_token_maps[s1]}
+    transitions_j = {(s1, states_to_token_maps[s1][j]) for s1 in states_to_token_maps if j in states_to_token_maps[s1]}
+    transitions_i, transitions_j = dict(transitions_i - transitions_j), dict(transitions_j - transitions_i)
+    for s1, s2 in transitions_i.items():
+        if s2 != transitions_j.get(s1):
+            states_to_token_maps[s1] = dict(states_to_token_maps[s1])
+            if s1 in transitions_j:
+                states_to_token_maps[s1][i] = transitions_j[s1]
+            else:
+                states_to_token_maps[s1].pop(i)
+            states_to_token_maps[s1][j] = s2
+    return states_to_token_maps
+
+
+def find_paths_with_transitions(states_to_token_maps: dict[int, dict[int, int]], transitions: list[int]) -> list[list[int]]:
+    possible_s0 = {s0 for s0 in states_to_token_maps if transitions[0] in states_to_token_maps[s0]}
+    possible_s1 = {s1 for s1 in states_to_token_maps if transitions[1] in states_to_token_maps[s1]} - possible_s0
+    starts = sorted(
+        s0 for s0 in possible_s0
+        if states_to_token_maps[s0][transitions[0]] in possible_s1
+    )
+    paths = [[start] for start in starts]
+    for path in paths:
+        for i in transitions:
+            if i in states_to_token_maps[path[-1]]:
+                path.append(states_to_token_maps[path[-1]][i])
+            else:
+                break
+    return [path for path in paths if len(path) == len(transitions) + 1]
+
+
+def replace_fields(fsm: RegexGuide, model: BaseModel, new_fields: list[str], tokenizer: PreTrainedTokenizerBase, make_infinite_loop: bool = False) -> RegexGuide:
+    assert len(new_fields) <= len(model.model_fields)
+    sttm = dict(fsm.states_to_token_maps)
+    encode = partial(tokenizer.encode, add_special_tokens=False)
+    quote = encode('"')[0]
+    # Let's replace the placeholder fields from the model in the finite state model by the new fields
+    for orig_field, new_field in zip(model.model_fields, new_fields):
+        orig_field_tokens = [encode(orig_field_char)[0] for orig_field_char in orig_field]
+        new_field_tokens = encode(new_field)
+        assert len(new_field_tokens) <= len(orig_field_tokens)
+        # Merge transitions until we have number of transitions = number of tokens in the field name
+        for k in reversed(range(len(new_field_tokens), len(orig_field_tokens))):
+            sttm = merge_successive_transitions(sttm, orig_field_tokens[k - 1], orig_field_tokens[k])
+        # Replace the token ids in the transitions with the ones of the new field name
+        for k in range(len(new_field_tokens)):
+            sttm = replace_transitions(sttm, orig_field_tokens[k], new_field_tokens[k])
+    if len(new_fields) < len(model.model_fields) or make_infinite_loop:
+        # Set the last field last state to generate less than the number of fields in the model
+        # We need to do this for every possible path
+        # e.g. multiple paths are used to count items when setting a min/max length
+        orig_last_field = list(model.model_fields)[-1]
+        new_last_field = new_fields[-1]
+        orig_last_field_paths = find_paths_with_transitions(sttm, [quote] + [encode(c)[0] for c in orig_last_field])
+        new_last_field_paths = find_paths_with_transitions(sttm, [quote] + encode(new_last_field))
+        if make_infinite_loop:  # this is a hack to loop on the same states over and over again
+            orig_last_field_paths = [orig_last_field_paths[0]] * len(orig_last_field_paths)
+        for orig_last_field_path, new_last_field_path in zip(
+            orig_last_field_paths,
+            new_last_field_paths
+        ):
+            orig_last_field_last_state = orig_last_field_path[-1]
+            new_last_field_second_last_state = new_last_field_path[-2]
+            sttm[new_last_field_second_last_state] = dict(sttm[new_last_field_second_last_state])
+            sttm[new_last_field_second_last_state][encode(new_last_field)[-1]] = orig_last_field_last_state
+    fsm = copy(fsm)
+    fsm.states_to_token_maps = sttm
+    return fsm

generate.py

@@ -0,0 +1,129 @@
+
+import json
+import logging
+import time
+from typing import Annotated, Iterator
+
+import ijson
+import outlines
+import torch
+from pydantic import BaseModel, StringConstraints, conlist, conset
+from outlines import generate, models
+from outlines.generate.api import SequenceGenerator
+from transformers import AutoTokenizer
+
+from fsm import replace_fields
+from samplers import PenalizedMultinomialSampler
+from utils import StringIteratorIO
+
+logger = logging.getLogger(__name__)
+
+
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    raise RuntimeError("couldn't find cuda or mps")
+
+logger.warning("Loading model...")
+model_id = "google/gemma-2b-it"
+# model_id = "Qwen/Qwen1.5-0.5B-Chat"
+model = models.transformers(model_id, device=device)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+sampler = PenalizedMultinomialSampler()
+empty_tokens = [token_id for token_id in range(tokenizer.vocab_size) if not tokenizer.decode([token_id]).strip()]
+sampler.set_max_repeats(empty_tokens, 1)
+
+# This Sample & Dataset models ztr just templated with placeholder fields
+
+class Sample(BaseModel):
+    # We use get_samples_generator() to replace the placeholder with the requested fields
+    ABCDabcd: str
+    EFGHefgh: str
+    IJKLijkl: str
+    MNOPmnop: str
+    QRSTqrst: str
+    # PS: don't use StringConstraints with max_length here since it creates a fsm that is too big
+
+
+class Dataset(BaseModel):
+    # We use get_samples_generator() to set the length to infinity
+    data: conlist(Sample, min_length=2, max_length=3)  # type: ignore
+
+
+samples_generator_template = generate.json(model, Dataset, sampler=sampler)
+
+class Columns(BaseModel):
+    columns: conset(Annotated[str, StringConstraints(pattern=r'[a-z0-9_]+')], min_length=2, max_length=len(Sample.model_fields))  # type: ignore
+
+columns_generator = generate.json(model, Columns, sampler=sampler)
+
+def get_samples_generator(new_fields: list[str]) -> SequenceGenerator:
+    fsm=samples_generator_template.fsm
+    fsm = replace_fields(  # replace the placeholder fields by the real fields
+        fsm=samples_generator_template.fsm,
+        model=Sample, 
+        new_fields=new_fields,
+        tokenizer=tokenizer,
+        make_infinite_loop=True  # to generate as many samples as we want
+    )
+    return SequenceGenerator(
+        fsm=fsm,
+        model=samples_generator_template.model,
+        sampler=samples_generator_template.sampler,
+        device=samples_generator_template.device
+    )
+
+
+@outlines.prompt
+def columns_prompt(filename: str):
+    """I would like to create a JSON file named {{ filename }}.json for a dataset of realistic data.
+    Give an example of column names / columns for this dataset to populate a SQL schema.
+    Please reply in JSON format and place the columns in a field named "columns".
+    """
+
+@outlines.prompt
+def samples_prommpt(filename: str, prompt: str, columns: str):
+    """I would like to create a JSON file named {{ filename }}.json for a dataset of realistic data.
+    Give an example of content using a JSON field named "data" with samples with columns {{ columns }}.
+    {{ prompt }}
+    """
+
+def stream_file(filename: str, prompt: str, columns: list[str], seed: int, size: int) -> Iterator[str]:
+    logger.warning(f"stream_response({filename=}, {prompt=}, {columns=})")
+    _start = time.time()
+    rng = torch.Generator(device=model.device)
+    rng.manual_seed(seed)
+    if not columns:
+
+        messages = [
+            {"role": "user", "content": columns_prompt(filename=filename)}
+        ]
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        logger.warning(f"stream_response({filename=}, {prompt=}, {columns=}) Generating columns...")
+        columns_generator_tokens = columns_generator.stream(text, rng=rng)
+        for column in ijson.items(StringIteratorIO(columns_generator_tokens), "columns.item", buf_size=16):
+            columns.append(column)
+        logger.warning(f"stream_response({filename=}, {prompt=}, {columns=}) Generating columns... DONE (total={time.time() - _start:.02f}s)")
+
+    logger.warning(f"stream_response({filename=}, {prompt=}, {columns=}) - Generating JSON regex guide...")
+    samples_generator = get_samples_generator(new_fields=columns)
+    logger.warning(f"stream_response({filename=}, {prompt=}, {columns=}) - Generating JSON regex guide... DONE (total={time.time() - _start:.02f}s)")
+    logger.warning(f"stream_response({filename=}, {prompt=}, {columns=}) - Generating samples...")
+    messages = [
+        {"role": "user", "content": samples_prommpt(filename=filename, prompt=prompt, columns="'" + "', '".join(columns) + "'")}
+    ]
+    text = tokenizer.apply_chat_template(
+        messages,
+        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)):
+        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

@@ -0,0 +1,77 @@
+import time
+
+import gradio as gr
+import io
+import pandas as pd
+import spaces
+
+@spaces.GPU(duration=120)
+def stream_output(filename: str):
+    if filename.endswith(".jsonl"):
+        filename = filename[:-len(".jsonl")]
+    from generate import stream_file
+    content = ""
+    size=3
+    start_time = time.time()
+    for i, chunk in enumerate(stream_file(
+        filename=filename,
+        prompt="",
+        columns=[],
+        seed=42,
+        size=size,
+    )):
+        content += chunk
+        df = pd.read_json(io.StringIO(content), lines=True)
+        state_msg = (
+            f"✅ Done generating {size} samples in {time.time() - start_time:.2f}s"
+            if i + 1 == size else
+            f"⚙️ Generating... [{i}/{size}]"
+        )
+        yield df, "```json\n" + content + "\n```", state_msg
+
+def test(filename: str):
+    if not filename.endswith(".jsonl"):
+        yield "❌ 404: File name must end with .jsonl", None, ""
+        return
+    
+    content = ""
+    size = 10
+    start_time = time.time()
+    for i in range(size):
+        content += f'{{"i": {i}, "filename": "{filename}"}}\n'
+        df = pd.read_json(io.StringIO(content), lines=True)
+        state_msg = (
+            f"✅ Done generating {size} samples in {time.time() - start_time:.2f}s"
+            if i + 1 == size else
+            f"⚙️ Generating... [{i}/{size}]"
+        )
+        yield df, "```json\n" + content + "\n```", state_msg
+        time.sleep(0.1)
+
+title = "LLM DataGen"
+description = "Generate and stream synthetic dataset files in JSON Lines format"
+examples = [
+    "movies_data.jsonl",
+    "common_first_names.jsonl",
+    "bad_amazon_reviews_on_defunct_products_that_people_hate.jsonl",
+    "dungeon_and_dragon_characters.jsonl"
+]
+
+with gr.Blocks() as demo:
+    gr.Markdown(f"# {title}")
+    gr.Markdown(description)
+    filename_comp = gr.Textbox(examples[0], placeholder=examples[0])
+    gr.Examples(examples, filename_comp)
+    generate_button = gr.Button("Generate dataset")
+    state_msg_comp = gr.Markdown("🔥 Ready to generate")
+    with gr.Tab("Dataset"):
+        dataframe_comp = gr.DataFrame()
+    with gr.Tab("File content"):
+        with gr.Blocks(fill_height=True):
+            with gr.Row():
+                file_content_comp = gr.Markdown()
+
+    generate_button.click(test, filename_comp, [dataframe_comp, file_content_comp, state_msg_comp])
+
+
+demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+spaces
+outlines==0.0.37
+transformers
+torch
+ijson
+pydantic

samplers.py

@@ -0,0 +1,72 @@
+import logging
+from typing import Tuple
+
+import torch
+from outlines.samplers import MultinomialSampler
+
+logger = logging.getLogger(__name__)
+
+class PenalizedMultinomialSampler(MultinomialSampler):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.penalized_tokens_group: list[torch.IntTensor] = []
+        self.max_repeats_per_token_group: list[int] = []
+        self.repeats_per_token_group: list[int] = []
+        self.token_id_to_tokens_groups: list[list[int]] = []
+    
+    def set_max_repeats(self, token_ids: list[int], max_repeats: int) -> None:
+        max_token_ids = max(token_ids)
+        if max_token_ids >= len(self.token_id_to_tokens_groups):
+            self.token_id_to_tokens_groups += [[] for _ in range(len(self.token_id_to_tokens_groups), max_token_ids + 1)]
+        for token_id in token_ids:
+            self.token_id_to_tokens_groups[token_id].append(len(self.penalized_tokens_group))
+        self.penalized_tokens_group.append(torch.tensor(token_ids, dtype=torch.int32))
+        self.max_repeats_per_token_group.append(max_repeats)
+        self.repeats_per_token_group.append(0)
+
+    def __call__(
+        self,
+        next_token_logits: torch.DoubleTensor,
+        sequence_weights: torch.DoubleTensor,
+        rng: torch.Generator,
+    ) -> Tuple[torch.DoubleTensor, torch.DoubleTensor, torch.DoubleTensor]:
+        """Call the multinomial sampler.
+
+        Parameters
+        ----------
+        next_token_logits
+            A tensor of shape ``(n_seqs, vocab_size,)`` that represents the
+            probability distribution of the next token over the vocabulary.
+        sequence_weights
+            A tensor of shape ``(n_seqs,)`` that represents the cumulative
+            weight of each sequence.
+        rng
+            A random number generator.
+
+        Returns
+        -------
+        A tuple with an array that contains the ids of the sampled tokens of
+        shape ``(n_seqs, 1)``, an array that contains the ancestors of each
+        sampled id of shape ``(n_seqs,)`` and an array that contains the updated
+        cumulative weights of each sequence of shape ``(n_seqs,)``.
+
+        """
+        if sequence_weights.min() == sequence_weights.max() == 0:
+            self.repeats_per_token_group = [0] * len(self.repeats_per_token_group)
+        else:
+            for penalized_tokens_group, max_repeats_per_token_group, repeats_per_token_group in zip(self.penalized_tokens_group, self.max_repeats_per_token_group, self.repeats_per_token_group):
+                if repeats_per_token_group >= max_repeats_per_token_group:
+                    penalty = torch.zeros_like(next_token_logits)
+                    penalty[:, penalized_tokens_group] = - torch.inf
+                    next_token_logits = next_token_logits + penalty
+        next_token_ids, ancestors, weights = super().__call__(
+            next_token_logits=next_token_logits,
+            sequence_weights=sequence_weights,
+            rng=rng
+        )
+        for next_token_id in next_token_ids.cpu():
+            if next_token_id < len(self.token_id_to_tokens_groups):
+                for token_group in self.token_id_to_tokens_groups[next_token_id]:
+                    self.repeats_per_token_group[token_group] += 1 
+        return next_token_ids, ancestors, weights

utils.py

@@ -0,0 +1,60 @@
+import io
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+class StringIteratorIO(io.TextIOBase):
+    """From: https://stackoverflow.com/a/12604375"""
+
+    def __init__(self, iter):
+        self._iter = iter
+        self._left = ''
+
+    def readable(self):
+        return True
+
+    def _read1(self, n=None):
+        while not self._left:
+            try:
+                self._left = next(self._iter)
+            except StopIteration:
+                break
+        ret = self._left[:n]
+        self._left = self._left[len(ret):]
+        return ret
+
+    def read(self, n=None):
+        buf = []
+        if n is None or n < 0:
+            while True:
+                m = self._read1()
+                if not m:
+                    break
+                buf.append(m)
+        else:
+            while n > 0:
+                m = self._read1(n)
+                if not m:
+                    break
+                n -= len(m)
+                buf.append(m)
+        return ''.join(buf)
+
+    def readline(self):
+        buf = []
+        while True:
+            i = self._left.find('\n')
+            if i == -1:
+                buf.append(self._left)
+                try:
+                    self._left = next(self._iter)
+                except StopIteration:
+                    self._left = ''
+                    break
+            else:
+                buf.append(self._left[:i+1])
+                self._left = self._left[i+1:]
+                break
+        return ''.join(buf)