cleaning

api.py

@@ -6,7 +6,7 @@ from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, StringConstraints
 from outlines import generate
 
-from generate import model, sampler, stream_file
+from generate import model, sampler, stream_jsonl_file
 
 logger = logging.getLogger(__name__)
 
@@ -22,7 +22,7 @@ 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(
+    for chunk in stream_jsonl_file(
         filename=filename,
         prompt=prompt,
         columns=columns,

gradio_app.py

@@ -1,25 +1,35 @@
 import time
+from urllib.parse import urlparse, parse_qs
 
 import gradio as gr
 import io
 import pandas as pd
 import spaces
 
-from generate import stream_file
+from generate import stream_jsonl_file
 
+MAX_SIZE = 20
+DEFAULT_SEED = 42
+DEFAULT_SIZE = 3
 
 @spaces.GPU(duration=120)
 def stream_output(filename: str):
-    if filename.endswith(".jsonl"):
-        filename = filename[:-len(".jsonl")]
+    parsed_filename = urlparse(filename)
+    filename = parsed_filename.path
+    params = parse_qs(parsed_filename.query)
+    prompt = params["prompt"][0] if "prompt" in params else ""
+    columns = [column.strip() for column in params["columns"][0].split(",") if column.strip()] if "columns" in params else []
+    size = int(params["size"][0]) if "size" in params else DEFAULT_SIZE
+    seed = int(params["seed"][0]) if "seed" in params else DEFAULT_SEED
+    if size > MAX_SIZE:
+        yield None, None, "Error: Maximum size is 20"
     content = ""
-    size=3
     start_time = time.time()
-    for i, chunk in enumerate(stream_file(
+    for i, chunk in enumerate(stream_jsonl_file(
         filename=filename,
-        prompt="",
-        columns=[],
-        seed=42,
+        prompt=prompt,
+        columns=columns,
+        seed=seed,
         size=size,
     )):
         content += chunk
@@ -31,32 +41,13 @@ def stream_output(filename: str):
         )
         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 + 1}/{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"
+    "bad_amazon_reviews_on_defunct_products_that_people_hate.jsonl",
+    "common_first_names.jsonl?columns=first_name,popularity&size=10",
 ]
 
 with gr.Blocks() as demo:
@@ -69,9 +60,7 @@ with gr.Blocks() as demo:
     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()
+        file_content_comp = gr.Markdown()
 
     generate_button.click(stream_output, filename_comp, [dataframe_comp, file_content_comp, state_msg_comp])