mostlyai-mock 0.1.18__py3-none-any.whl → 0.2.1__py3-none-any.whl

mostlyai/mock/__init__.py

@@ -15,4 +15,4 @@
 from mostlyai.mock.core import sample
 
 __all__ = ["sample"]
-__version__ = "0.1.18"  # Do not set this manually. Use poetry version [params].
+__version__ = "0.2.1"  # Do not set this manually. Use poetry version [params].

mostlyai/mock/core.py

@@ -18,8 +18,9 @@ import asyncio
 import concurrent.futures
 import json
 import math
+import time
 from collections import deque
-from collections.abc import AsyncGenerator, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from enum import Enum
 from io import StringIO
 from typing import Any, Literal
@@ -29,7 +30,6 @@ import litellm
 import pandas as pd
 import tenacity
 from pydantic import BaseModel, Field, RootModel, create_model, field_validator, model_validator
-from tqdm.asyncio import tqdm
 
 litellm.suppress_debug_info = True
 
@@ -249,8 +249,26 @@ async def _sample_table(
     n_workers: int,
     llm_config: LLMConfig,
     config: MockConfig,
-    progress_callback: Callable | None = None,
+    progress_callback: Callable[..., Awaitable[None]] | None = None,
 ) -> pd.DataFrame:
+    # provide a default progress callback if none is provided
+    if progress_callback is None:
+
+        async def default_progress_callback(**kwargs):
+            percentage = (kwargs["progress"] / kwargs["total"]) * 100 if kwargs["total"] > 0 else 0
+            rows_per_second = kwargs["rows"] / kwargs["elapsed_time"] if kwargs["elapsed_time"] > 0 else 0
+            message = (
+                f"Generating table `{kwargs['table']}`".ljust(40)
+                + f": {percentage:3.0f}%, {kwargs['rows']} rows, {kwargs['elapsed_time']:.0f}s, {rows_per_second:.1f} rows/s"
+            )
+            is_final = kwargs["progress"] >= kwargs["total"]
+            if is_final:
+                print(f"\r{message}")  # final update with newline
+            else:
+                print(f"\r{message}", end="", flush=True)  # in-progress update
+
+        progress_callback = default_progress_callback
+
     table_rows_generator = _create_table_rows_generator(
         name=name,
         prompt=prompt,
@@ -265,7 +283,6 @@ async def _sample_table(
         llm_config=llm_config,
         progress_callback=progress_callback,
     )
-    table_rows_generator = tqdm(table_rows_generator, desc=f"Generating rows for table `{name}`".ljust(45))
     table_df = await _convert_table_rows_generator_to_df(
         table_rows_generator=table_rows_generator,
         columns=columns,
@@ -805,7 +822,7 @@ async def _create_table_rows_generator(
     non_context_size: int | None,
     n_workers: int,
     llm_config: LLMConfig,
-    progress_callback: Callable | None = None,
+    progress_callback: Callable[..., Awaitable[None]] | None = None,
 ) -> AsyncGenerator[dict]:
     batch_size = 20  # generate 20 root table rows at a time
 
@@ -863,6 +880,10 @@ async def _create_table_rows_generator(
         # +2 because LLM may not always count the rows correctly
         batch_sizes[-1] = sample_size - sum(batch_sizes[:-1]) + 2
 
+    # emit initial progress message right away
+    if progress_callback:
+        await progress_callback(table=name, progress=0, total=n_total_batches, rows=0, elapsed_time=0)
+
     # initialize queues for async communication
     batch_queue = asyncio.PriorityQueue()
     result_queue = asyncio.Queue()
@@ -927,6 +948,8 @@ async def _create_table_rows_generator(
 
     n_completed_batches = 0
     n_yielded_sequences = 0
+    n_generated_rows = 0
+    table_start_time = time.time()
     while n_yielded_sequences < sample_size:
         if n_completed_batches >= n_total_batches:
             assert context_data is None, "n_total_batches is fixed for linked tables"
@@ -950,6 +973,7 @@ async def _create_table_rows_generator(
         rows = result
         for row_idx, row in enumerate(rows):
             yield (batch_idx, row)
+            n_generated_rows += 1
             if context_batches is None or row_idx == len(rows) - 1:
                 # in case of flat table, each row is considered a single sequence
                 # in case of linked table, all rows are considered a single sequence
@@ -959,10 +983,13 @@ async def _create_table_rows_generator(
                 break
         n_completed_batches += 1
         if progress_callback:
+            elapsed_time = time.time() - table_start_time
             await progress_callback(
+                table=name,
                 progress=n_completed_batches,
                 total=n_total_batches,
-                message=f"Generating rows for table `{name}`: {n_completed_batches}/{n_total_batches}",
+                rows=n_generated_rows,
+                elapsed_time=round(elapsed_time, 2),
             )
         result_queue.task_done()
 
@@ -1232,7 +1259,7 @@ async def _sample_common(
     top_p: float = 0.95,
     n_workers: int = 10,
     return_type: Literal["auto", "dict"] = "auto",
-    progress_callback: Callable | None = None,
+    progress_callback: Callable[..., Awaitable[None]] | None = None,
 ):
     tables: dict[str, TableConfig] = _harmonize_tables(tables, existing_data)
     config = MockConfig(tables)
@@ -1290,6 +1317,7 @@ def sample(
     top_p: float = 0.95,
     n_workers: int = 10,
     return_type: Literal["auto", "dict"] = "auto",
+    progress_callback: Callable[..., Awaitable[None]] | None = None,
 ) -> pd.DataFrame | dict[str, pd.DataFrame]:
     """
     Generate synthetic data from scratch or enrich existing data with new columns.
@@ -1305,6 +1333,7 @@ def sample(
     Args:
         tables (dict[str, dict]): The table specifications to generate mock data for. See examples for usage.
             Note: Avoid using double quotes (`"`) and other special characters in column names.
+            Available dtypes: `string`, `integer`, `float`, `category`, `boolean`, `date`, `datetime`.
         sample_size (int | dict[str, int]): The number of rows to generate for each subject table.
             If a single integer is provided, the same number of rows will be generated for each subject table.
             If a dictionary is provided, the number of rows to generate for each subject table can be specified individually.
@@ -1329,6 +1358,11 @@ def sample(
         n_workers (int): The number of concurrent workers making the LLM calls. Default is 10. The value is clamped to the range [1, 10].
             If n_workers is 1, the generation of batches becomes sequential and certain features for better data consistency are enabled.
         return_type (Literal["auto", "dict"]): The format of the returned data. Default is "auto".
+        progress_callback (Callable | None): Optional callback function to track progress during data generation.
+            If not provided, a default progress callback will display progress messages in the format:
+            "Generating table `table_name`: X%, Y rows, Zs, W.X rows/s"
+            The callback receives keyword arguments including: table, progress, total,
+            rows, and elapsed_time. Default is None.
 
     Returns:
         - pd.DataFrame: A single DataFrame containing the generated mock data, if only one table is provided.
@@ -1506,6 +1540,26 @@ def sample(
     df_customers = data["customers"]
     df_orders = data["orders"]
     ```
+
+    Example of using a custom progress callback to provide progress in JSON format:
+    ```python
+    from mostlyai import mock
+    import asyncio
+    import json
+
+    async def custom_progress_callback(**kwargs):
+        msg = f"\r{json.dumps(kwargs)}"
+        if kwargs["progress"] < kwargs["total"]:
+            print(msg, end="", flush=True)
+        else:
+            print(msg)
+
+    df = mock.sample(
+        tables=tables,
+        sample_size=10,
+        progress_callback=custom_progress_callback
+    )
+    ```
     """
 
     def sample_common_sync(*args, **kwargs) -> pd.DataFrame | dict[str, pd.DataFrame]:
@@ -1523,7 +1577,7 @@ def sample(
             top_p=top_p,
             n_workers=n_workers,
             return_type=return_type,
-            progress_callback=None,
+            progress_callback=progress_callback,
         )
         return future.result()
 
@@ -1539,7 +1593,7 @@ async def _asample(
     top_p: float = 0.95,
     n_workers: int = 10,
     return_type: Literal["auto", "dict"] = "auto",
-    progress_callback: Callable | None = None,
+    progress_callback: Callable[..., Awaitable[None]] | None = None,
 ) -> pd.DataFrame | dict[str, pd.DataFrame]:
     return await _sample_common(
         tables=tables,

{mostlyai_mock-0.1.18.dist-info → mostlyai_mock-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-mock
-Version: 0.1.18
+Version: 0.2.1
 Summary: Synthetic Mock Data
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-mock
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-mock
@@ -47,7 +47,7 @@ Use LLMs to generate any Tabular Data towards your needs. Create from scratch, e
 * A light-weight python client for prompting LLMs for mixed-type tabular data.
 * Select from a wide range of LLM endpoints and LLM models.
 * Supports single-table as well as multi-table scenarios.
-* Supports variety of data types: `string`, `categorical`, `integer`, `float`, `boolean`, `date`, and `datetime`.
+* Supports variety of data types: `string`, `integer`, `float`, `category`, `boolean`, `date`, and `datetime`.
 * Specify context, distributions and rules via dataset-, table- or column-level prompts.
 * Create from scratch or enrich existing datasets with new columns and/or rows.
 * Tailor the diversity and realism of your generated data via temperature and top_p.

mostlyai_mock-0.2.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+mostlyai/mock/__init__.py,sha256=JUs456bSY9qLCZk_m-QjtcFDf4e3c9DkzdQpBafJZKA,714
+mostlyai/mock/core.py,sha256=TsBPnV068Dmgy52sWKDKe055_LWQtWy9YO4Jx5bRVLk,69227
+mostlyai/mock/mcp_server.py,sha256=uDLg0SeMPV2VZhXviM-F769W0xlmhGwlmQiQhY0Q-Ik,2365
+mostlyai_mock-0.2.1.dist-info/METADATA,sha256=eSwEcdV3RwSUw6ciEEjeBAzV99E_mf1wj9XGXjapTak,14253
+mostlyai_mock-0.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mostlyai_mock-0.2.1.dist-info/entry_points.txt,sha256=XDbppUIAaCWW0nresVep8zb71pkzZuFA16jCBHq8CU8,61
+mostlyai_mock-0.2.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mostlyai_mock-0.2.1.dist-info/RECORD,,

mostlyai_mock-0.1.18.dist-info/RECORD

@@ -1,8 +0,0 @@
-mostlyai/mock/__init__.py,sha256=UKmnKlQ7fZVvB0ckh9_nXjojAE0JGa2Kd2mT0Ci8cDU,715
-mostlyai/mock/core.py,sha256=oGSpIXINL7R1X7ZN5dtdwItaPXDD0mGvkakA0CEzmwI,66880
-mostlyai/mock/mcp_server.py,sha256=uDLg0SeMPV2VZhXviM-F769W0xlmhGwlmQiQhY0Q-Ik,2365
-mostlyai_mock-0.1.18.dist-info/METADATA,sha256=EmLjpo-D-wJefswHIMk3TCK9TvzLML-3Sjo0OEi9qAI,14257
-mostlyai_mock-0.1.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mostlyai_mock-0.1.18.dist-info/entry_points.txt,sha256=XDbppUIAaCWW0nresVep8zb71pkzZuFA16jCBHq8CU8,61
-mostlyai_mock-0.1.18.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mostlyai_mock-0.1.18.dist-info/RECORD,,