From 0a5705d03579e17e8de30993fae7da2cfec5b94a Mon Sep 17 00:00:00 2001 From: "Massimo R. Scamarcia" Date: Thu, 27 Aug 2026 19:10:28 +0200 Subject: [PATCH] Add topics and inspect commands, SFT-ready chat, and generated-data ignores Single squashed commit for the topics-command MR, covering the branch's full contribution from its initial commit through the latest changes. - topics: LLM-powered conversation topic generation (structured output, --append incremental growth, category-driven diversity) - chat: multi-turn SFT-ready conversation generation with reasoning control, --max-tokens, and strictly alternating roles - Incremental save, resume, and retry for chat and topics batch runs - inspect: stream any completionist dataset, show local parquet totals, and report elapsed time instead of a loading bar - Ignore generated dataset outputs (*.jsonl, docs/examples/*, .qwen) via .gitignore and .dockerignore; bump version to 0.1.2 --- .dockerignore | 12 ++ .gitignore | 7 + README.md | 157 +++++++++++++++- completionist/cli.py | 4 + completionist/commands/chat.py | 162 +++++++++++----- completionist/commands/inspect.py | 255 +++++++++++++++++++++++++ completionist/commands/topics.py | 301 ++++++++++++++++++++++++++++++ completionist/dataset_stream.py | 135 ++++++++++++++ completionist/dataset_view.py | 136 ++++++++++++++ completionist/pager.py | 107 +++++++++++ completionist/processing.py | 5 +- pyproject.toml | 8 +- tests/test_chat.py | 225 ++++++++++++++++++++++ tests/test_dataset_stream.py | 192 +++++++++++++++++++ tests/test_dataset_view.py | 152 +++++++++++++++ tests/test_inspect.py | 157 ++++++++++++++++ tests/test_pager.py | 91 +++++++++ tests/test_processing.py | 76 ++++++++ tests/test_topics.py | 71 +++++++ uv.lock | 149 ++++++++++++++- 20 files changed, 2349 insertions(+), 53 deletions(-) create mode 100644 .dockerignore create mode 100644 completionist/commands/inspect.py create mode 100644 completionist/commands/topics.py create mode 100644 completionist/dataset_stream.py create mode 100644 completionist/dataset_view.py create mode 100644 completionist/pager.py create mode 100644 tests/test_chat.py create mode 100644 tests/test_dataset_stream.py create mode 100644 tests/test_dataset_view.py create mode 100644 tests/test_inspect.py create mode 100644 tests/test_pager.py create mode 100644 tests/test_processing.py create mode 100644 tests/test_topics.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..5fd6ca0 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,12 @@ +.git +.github +.qwen +.venv +__pycache__ +*.py[co] +.pytest_cache +.ruff_cache +scratch/ +datasets/ +*.jsonl +docs/examples/*/ diff --git a/.gitignore b/.gitignore index 938ada2..942534a 100644 --- a/.gitignore +++ b/.gitignore @@ -12,3 +12,10 @@ wheels/ # Local scripts and generated data scratch/ datasets/ + +# Generated dataset outputs (batch runs, not for distribution) +*.jsonl +docs/examples/*/ + +# Local assistant state +.qwen/ diff --git a/README.md b/README.md index 70442b1..b5c7bba 100644 --- a/README.md +++ b/README.md @@ -87,8 +87,159 @@ uv run python3 -m completionist build \ --schema completionist.default_schema.SchemaWithReasoning ``` -## Future Development +## Generating Multi-Turn Conversations with chat + +The `chat` command generates multi-turn conversation datasets from a list of topics. Each topic seeds one or more realistic dialogues between a `user` and an `assistant`, produced as structured JSON via outlines (no `--schema` flag needed). + +``` +uv run python3 -m completionist chat \ + --topics-file topics.txt \ + --num-conversations 3 \ + --min-turns 4 \ + --max-turns 6 \ + --model-name Gemma-4-E2B-it-GGUF \ + --api-url http://localhost:11434/v1 \ + --output-file chat_dataset.parquet +``` + +This generates 3 conversations for each topic in `topics.txt` (one topic per line), each with 4–6 messages. + +### Key options + +| Option | Description | +|--------|-------------| +| `--topics-file` | Path to a text file with one topic per line. **Required.** | +| `--num-conversations` | Number of conversations to generate per topic. **Required.** | +| `--min-turns` / `--max-turns` | Message count range. Defaults `4`/`6`. **Must be even** so the conversation ends on an `assistant` turn (user-first alternation). | +| `--output-file` | Output path (`.parquet` or `.jsonl`). **Required.** | +| `--model-name` | Model to use. **Required.** | +| `--api-url` | OpenAI-compatible endpoint. Defaults to Ollama. | +| `--system-prompt` / `--system-prompt-file` | Override the built-in conversational system prompt. | +| `--user-prompt-template` | Override the user-prompt template (`{topic}` and `{num_turns}` are substituted). | +| `--workers` | Concurrent requests. Defaults to `4`; set to `1` for single-request servers like llama.cpp. | +| `--temperature` / `--top-p` | Sampling controls. Defaults `0.7` / `0.95`. | +| `--max-tokens` | Token budget per conversation. Defaults to `4096`. | +| `--reasoning` | Reasoning mode: `on`, `off`, or `auto`. | +| `--push-to-hub` / `--hf-repo-id` | Push the result to the Hugging Face Hub. | + +The output schema has a `topic` field and a `messages` list of `{role, content}` objects, with roles strictly alternating `user` → `assistant` starting with `user`. + +## Generating Topics with topics + +The `topics` command generates conversation topic lists with an LLM — useful for seeding the `chat` command. It requests topics in batches and deduplicates automatically across batches. + +``` +uv run python3 -m completionist topics \ + --num-topics 100 \ + --categories "technology, philosophy, science, ethics" \ + --model-name Gemma-4-E2B-it-GGUF \ + --api-url http://localhost:11434/v1 \ + --output topics.txt +``` + +Output format is determined by the file extension: + +- `.txt` — one topic per line +- `.jsonl` — one `{"topic": "..."}` object per line + +### Key options + +| Option | Description | +|--------|-------------| +| `--output` / `-o` | Output path. `.txt` or `.jsonl` extension selects the format. **Required.** | +| `--num-topics` / `-n` | Number of unique topics to generate. **Required.** | +| `--model-name` | Model to use. **Required.** | +| `--categories` | Comma-separated category hints to guide diversity. **Required.** | +| `--api-url` | OpenAI-compatible endpoint. Defaults to Ollama. | +| `--batch-size` | Topics per API call. Defaults to `50`. | +| `--temperature` / `--top-p` | Sampling controls. Defaults `0.9` / `0.95`. | +| `--max-tokens` | Token budget per batch. Defaults to `4096`. | +| `--system-prompt` / `--system-prompt-file` | Override the built-in topic-curation prompt. | +| `--reasoning` | Reasoning mode: `on`, `off`, or `auto`. | +| `--max-batches` | Upper bound on API calls before stopping. Defaults to `20`. | +| `--append` | Append to an existing file; `--num-topics` becomes how many *new* topics to add on top of what is already there. | + +Without `--append`, the output file is overwritten. With `--append`, existing topics are loaded first and only new unique topics are generated, so repeated runs accumulate a non-overlapping pool. + +## Translating Datasets with translate + +The `translate` command translates one or more text fields of a dataset using an LLM, with optional Redis caching to skip already-translated strings across runs. -This tool's functionality will be expanded in the near future to support different tasks. +``` +uv run python3 -m completionist translate \ + --dataset-name ethicalabs/kurtis-v2-sft-mix-tiny \ + --input-field prompt \ + --input-field completion \ + --input-field reasoning \ + --source-lang English \ + --target-lang Italian \ + --model-name Gemma-4-E4B-it-GGUF \ + --api-url http://localhost:11434/v1 \ + --output-file translated.parquet +``` + +Repeat `--input-field` once per field to translate. For each input field the output adds two columns: `source_{field}` and `translated_{field}` (e.g. `source_prompt` / `translated_prompt`). + +### Key options + +| Option | Description | +|--------|-------------| +| `--dataset-name` | Hugging Face dataset name. **Required.** | +| `--input-field` | Field to translate. Repeatable for multiple fields. **Required.** | +| `--source-lang` / `--target-lang` | Language pair. **Required.** | +| `--output-file` | Output path. **Required.** | +| `--model-name` | Translation model. Defaults to `translategemma-4b-it-GGUF-Q4_K_M`. | +| `--api-url` | OpenAI-compatible endpoint. Defaults to Ollama. | +| `--system-prompt` / `--system-prompt-file` | Override the generated translation prompt (mutually exclusive). | +| `--limit` | Process only the first N samples. | +| `--shuffle` | Shuffle the dataset before processing. | +| `--workers` | Concurrent requests. Defaults to `4`. | +| `--temperature` / `--top-p` / `--max-tokens` | Sampling controls. Defaults `0.7` / `0.95` / `4096`. | +| `--cache-url` | Redis URL (e.g. `redis://localhost:6379`) to cache translations. Requires `pip install redis` and a running Redis server. | +| `--push-to-hub` / `--hf-repo-id` | Push the result to the Hugging Face Hub. | + +## Reading Datasets with inspect + +The `inspect` command reads any dataset completionist produces and auto-detects its shape — no flags needed to pick a view. It opens an interactive pager, or falls back to a plain-text dump when the output is piped. + +``` +uv run python3 -m completionist inspect datasets/chat_sft_merged.parquet +uv run python3 -m completionist inspect ethical-01.jsonl +uv run python3 -m completionist inspect ethicalabs/some-dataset --split train +``` + +### Detected formats + +| Format | Detection | Tabs | +|--------|-----------|------| +| Conversation | `messages` (list of `{role, content}`) | `conversation`, `raw` | +| Prompt/completion | `prompt` + `completion` columns | `prompt`, `completion`, `reasoning` (if present), `raw` | +| Tools | `messages` with `tool_calls`/`tool` turns, or a `tools` field | `conversation`, `tools`, `raw` | +| Fallback | anything else | `raw` (pretty-printed JSON of every field) | + +The `raw` tab is always present, so any dataset is inspectable. + +### Keys + +- `←` / `→` — previous / next sample +- `Tab` / `Shift-Tab` — cycle between tabs +- `s` — cycle between splits (when the dataset has several) +- `↑` / `↓` — scroll +- `PgUp` / `PgDn` — page up / down +- `Home` / `End` — jump to top / bottom +- `q` / `Esc` — quit + +### Options + +| Option | Description | +|--------|-------------| +| `dataset` | Local `.parquet`/`.jsonl` file, local directory (HF split layout), or a Hugging Face dataset name. **Positional.** | +| `--split` | Split to load. Auto-detected (`train` → `test` → `validation` → first) if omitted. | +| `--limit` | Load at most N samples. | +| `--index` | Skip this many samples before reading. | + +Datasets are read via streaming — remote and local alike — with a bounded in-memory window; older samples spill to a `/tmp` cache so backward navigation works without holding the whole dataset in RAM. + +## Future Development -The plan is to add subcommands to handle various scenarios, such as translating text and post-processing. +This tool's functionality will be expanded in the near future to support additional dataset-processing tasks and post-processing utilities. diff --git a/completionist/cli.py b/completionist/cli.py index 812b8c8..b3cd87f 100644 --- a/completionist/cli.py +++ b/completionist/cli.py @@ -5,6 +5,8 @@ from completionist.commands.complete import complete_cmd # from completionist.commands.compose import compose_cmd +from completionist.commands.inspect import inspect_cmd +from completionist.commands.topics import topics_cmd from completionist.commands.translate import translate_cmd @@ -17,4 +19,6 @@ def entry_point(): entry_point.add_command(chat_cmd) entry_point.add_command(complete_cmd) # entry_point.add_command(compose_cmd) +entry_point.add_command(inspect_cmd) +entry_point.add_command(topics_cmd) entry_point.add_command(translate_cmd) diff --git a/completionist/commands/chat.py b/completionist/commands/chat.py index 572d8a6..6f70fb6 100644 --- a/completionist/commands/chat.py +++ b/completionist/commands/chat.py @@ -1,11 +1,14 @@ import os import sys +import json import random import traceback +import time import click from pydantic import BaseModel, Field from huggingface_hub import get_token +from datasets import Dataset from completionist.processing import process_samples_with_executor from completionist.dataset_io import save_and_push_dataset @@ -33,8 +36,14 @@ class ChatConversation(BaseModel): "Generate a realistic, engaging multi-turn conversation on the given topic. " "The conversation must have exactly {num_turns} messages, alternating between " "'user' and 'assistant' roles, starting with 'user'. " - "Vary the tone and depth — some turns short and casual, others longer and thoughtful. " - "Return the result as a JSON object matching the schema." + "Make it feel like two interesting people talking, not a Q&A or a lecture.\n" + "- The assistant should occasionally disagree, push back, ask a question, or admit " + "they don't know. Avoid wrapping up neatly — real conversations are messy.\n" + "- Vary register: humor, frustration, curiosity, doubt. Not every turn should be " + "polished and academic.\n" + "- The user should have their own perspective, not just ask setup questions. " + "Let them challenge the assistant, change their mind, or go off on a tangent.\n" + "- Return the result as a JSON object matching the schema." ) DEFAULT_USER_PROMPT_TEMPLATE = ( @@ -50,44 +59,56 @@ class ChatConversation(BaseModel): def chat_task_handler(topic: str, llm_config: dict): """Task handler for generating a single multi-turn conversation for a topic.""" - num_turns = random.randint(llm_config["min_turns"], llm_config["max_turns"]) - - user_prompt = llm_config["user_prompt_template"].format( - topic=topic, num_turns=num_turns + num_turns = random.randrange( + llm_config["min_turns"], llm_config["max_turns"] + 1, 2 ) - system_prompt = llm_config["system_prompt"].format(num_turns=num_turns) - - try: - result = get_completion( - prompt=user_prompt, - model_name=llm_config["model_name"], - api_url=llm_config["api_url"], - system_prompt=system_prompt, - hf_api_token=llm_config["hf_api_token"], - openai_api_token=llm_config["openai_api_token"], - pydantic_schema=ChatConversation, - temperature=llm_config["generation_config"]["temperature"], - top_p=llm_config["generation_config"]["top_p"], - max_tokens=llm_config.get("max_tokens", 2048), - reasoning=llm_config.get("reasoning"), - ) - - if result is None: - return None - if isinstance(result, ChatConversation): - return result.model_dump() - # outlines returned a raw string — attempt JSON parse as fallback - import json - raw = json.loads(result) - return ChatConversation(**raw).model_dump() - - except Exception: - print( - f"\nWarning: Failed to generate conversation for topic '{topic}'. " - f"Reason: {traceback.format_exc()}" - ) - return None + user_prompt = ( + llm_config["user_prompt_template"] + .replace("{topic}", topic) + .replace("{num_turns}", str(num_turns)) + ) + system_prompt = llm_config["system_prompt"].replace("{num_turns}", str(num_turns)) + + max_retries = llm_config.get("retries", 3) + for attempt in range(1, max_retries + 1): + try: + result = get_completion( + prompt=user_prompt, + model_name=llm_config["model_name"], + api_url=llm_config["api_url"], + system_prompt=system_prompt, + hf_api_token=llm_config["hf_api_token"], + openai_api_token=llm_config["openai_api_token"], + pydantic_schema=ChatConversation, + temperature=llm_config["generation_config"]["temperature"], + top_p=llm_config["generation_config"]["top_p"], + max_tokens=llm_config.get("max_tokens", 2048), + reasoning=llm_config.get("reasoning"), + ) + + if result is None: + raise RuntimeError("completion returned no result") + if isinstance(result, ChatConversation): + return result.model_dump() + # outlines returned a raw string — attempt JSON parse as fallback + return ChatConversation(**json.loads(result)).model_dump() + + except Exception as exc: + if attempt >= max_retries: + print( + f"\nWarning: Failed to generate conversation for topic '{topic}' " + f"after {max_retries} attempts. Reason: {traceback.format_exc()}" + ) + return None + delay = 2 ** (attempt - 1) + print( + f" Retrying topic '{topic}' (attempt {attempt}/{max_retries} failed): " + f"{exc} — waiting {delay}s" + ) + time.sleep(delay) + + return None # --- CLI command --- @@ -109,15 +130,15 @@ def chat_task_handler(topic: str, llm_config: dict): @click.option( "--min-turns", type=int, - default=3, - help="Minimum number of messages per conversation.", + default=4, + help="Minimum number of messages per conversation. Must be even.", show_default=True, ) @click.option( "--max-turns", type=int, default=6, - help="Maximum number of messages per conversation.", + help="Maximum number of messages per conversation. Must be even.", show_default=True, ) @click.option( @@ -189,6 +210,13 @@ def chat_task_handler(topic: str, llm_config: dict): help="Maximum tokens per conversation. Increase for longer chats.", show_default=True, ) +@click.option( + "--retries", + type=int, + default=3, + help="Maximum attempts per conversation on failure.", + show_default=True, +) def chat_cmd( topics_file, num_conversations, @@ -207,6 +235,7 @@ def chat_cmd( top_p, reasoning, max_tokens, + retries, ): """ Generate multi-turn conversation datasets from a list of topics. @@ -219,10 +248,18 @@ def chat_cmd( print("Error: --hf-repo-id is required when --push-to-hub is used.") sys.exit(1) + if min_turns % 2 != 0 or max_turns % 2 != 0: + print("Error: --min-turns and --max-turns must be even numbers.") + sys.exit(1) + if min_turns > max_turns: print("Error: --min-turns cannot be greater than --max-turns.") sys.exit(1) + if retries < 1: + print("Error: --retries must be at least 1.") + sys.exit(1) + # Resolve prompts: file > inline > built-in default system_prompt_content = read_file_content(system_prompt_file) or system_prompt if not system_prompt_content: @@ -255,6 +292,7 @@ def chat_cmd( "openai_api_token": openai_api_token, "reasoning": reasoning, "max_tokens": max_tokens, + "retries": retries, } # Build task list: num_conversations per topic @@ -262,21 +300,55 @@ def chat_cmd( for topic in topics: tasks.extend([topic] * num_conversations) + # Resume from an existing output file if present. + existing = [] + resume_idx = 0 + if os.path.exists(output_file): + try: + if os.path.splitext(output_file)[1].lower() == ".jsonl": + existing = Dataset.from_json(output_file).to_list() + else: + existing = Dataset.from_parquet(output_file).to_list() + resume_idx = len(existing) + print( + f"Resuming from {output_file}: {resume_idx} conversations already done." + ) + except Exception as e: + print(f"Could not load {output_file} ({e}). Starting fresh.") + + remaining = tasks[resume_idx:] + + if not remaining: + print(f"All {len(tasks)} conversations already present in {output_file}.") + return + print( - f"Generating {len(tasks)} conversations ({num_conversations} per topic) " + f"Generating {len(remaining)} conversations ({num_conversations} per topic) " f"across {len(topics)} topics with {workers} workers..." ) + def save_progress(completions): + # Checkpoint partial results locally so interrupted runs keep their output. + save_and_push_dataset( + completions=existing + completions, + output_file=output_file, + push_to_hub=False, + hf_repo_id=None, + hf_api_token=None, + ) + generated = process_samples_with_executor( - dataset_to_process=tasks, + dataset_to_process=remaining, workers=workers, - resume_idx=0, + resume_idx=resume_idx, task_handler=chat_task_handler, llm_config=llm_config, + save_callback=save_progress, + save_every=25, ) save_and_push_dataset( - completions=generated, + completions=existing + generated, output_file=output_file, push_to_hub=push_to_hub, hf_repo_id=hf_repo_id, diff --git a/completionist/commands/inspect.py b/completionist/commands/inspect.py new file mode 100644 index 0000000..c0f2ff1 --- /dev/null +++ b/completionist/commands/inspect.py @@ -0,0 +1,255 @@ +"""Read a completionist dataset interactively. + +Usage: + completionist inspect + +The format is auto-detected per sample (conversation, prompt/completion, +tools, or a raw JSON fallback). Local files, local directories following the +HF split layout, and remote datasets are all read via streaming. + +Keys: + left / right previous / next sample + tab / shift-tab cycle view + s cycle split (when the dataset has several) + up / down scroll + PgUp / PgDn page up / down + Home / End jump to top / bottom + q / Esc quit +""" + +import curses +import os +import sys +import textwrap +import time + +import click + +from completionist.dataset_stream import DatasetStream +from completionist.dataset_view import ( + ROLE_PREFIX, + detect_format, + render_sample_text, + row_to_tabs, + title, +) +from completionist.pager import Pager +from completionist.utils import handle_error + + +def _wrap_segments(segments, width): + """Wrap (style, text) segments into display lines for the given width.""" + width = max(1, width) + lines = [] + for style, text in segments: + if style == "raw": + lines.extend((style, ln) for ln in text.splitlines()) + continue + prefix = ROLE_PREFIX.get(style, "") + if prefix: + wrapped = textwrap.wrap(text, width=max(1, width - len(prefix))) or [""] + lines.append((style, prefix + wrapped[0])) + for cont in wrapped[1:]: + lines.append((style, " " * len(prefix) + cont)) + else: + for chunk in textwrap.wrap(text, width=width) or [""]: + lines.append((style, chunk)) + return lines + + +def _tui(stdscr, source, pager): + curses.curs_set(0) + stdscr.keypad(True) + if curses.has_colors(): + curses.start_color() + curses.use_default_colors() + curses.init_pair(1, curses.COLOR_CYAN, -1) # user + curses.init_pair(2, curses.COLOR_GREEN, -1) # assistant + curses.init_pair(3, curses.COLOR_YELLOW, -1) # topic + curses.init_pair(4, curses.COLOR_BLACK, curses.COLOR_WHITE) # bars + curses.init_pair(5, curses.COLOR_MAGENTA, -1) # tool + + attrs = { + "user": curses.color_pair(1), + "assistant": curses.color_pair(2), + "topic": curses.color_pair(3) | curses.A_BOLD, + "tool": curses.color_pair(5), + "system": curses.A_BOLD, + "raw": curses.A_NORMAL, + "plain": curses.A_NORMAL, + } + + tab_idx = 0 + scroll = 0 + + while True: + row = pager.current() + fmt = detect_format(row) + tabs = row_to_tabs(row, fmt) + tab_idx %= len(tabs) + _, segments = tabs[tab_idx] + + stdscr.erase() + H, W = stdscr.getmaxyx() + body_h = max(1, H - 3) + + total = pager.total + split_label = f" [{source.split}]" if source.split else "" + header = ( + f" [{pager.cursor + 1}/{total if total is not None else '?'}]" + f"{split_label} {title(row, fmt)}" + ) + try: + stdscr.addnstr(0, 0, header, W, curses.color_pair(4) | curses.A_BOLD) + except curses.error: + pass + + tabbar = " ".join( + f"[{name}]" if i == tab_idx else name for i, (name, _) in enumerate(tabs) + ) + try: + stdscr.addnstr(1, 0, tabbar, W, curses.color_pair(4)) + except curses.error: + pass + + lines = _wrap_segments(segments, W) + max_scroll = max(0, len(lines) - body_h) + scroll = max(0, min(scroll, max_scroll)) + for i in range(body_h): + li = scroll + i + if li >= len(lines): + break + style, text = lines[li] + try: + stdscr.addnstr(2 + i, 0, text, W, attrs.get(style, curses.A_NORMAL)) + except curses.error: + pass + + keys = ["←/→ sample", "Tab/Shift-Tab view"] + if len(source.splits) > 1: + keys.append("s split") + keys += ["↑/↓ scroll", "PgUp/PgDn page", "q quit"] + footer = " " + " ".join(keys) + " " + try: + stdscr.addnstr( + H - 1, + 0, + footer.ljust(W)[:W], + W, + curses.color_pair(4) | curses.A_REVERSE, + ) + except curses.error: + pass + + stdscr.refresh() + + key = stdscr.getch() + if key in (ord("q"), 27): + break + elif key == curses.KEY_LEFT: + pager.prev() + scroll = 0 + elif key == curses.KEY_RIGHT: + pager.next() + scroll = 0 + elif key == ord("s"): + if source.switch_split(): + try: + first = source.next() + except StopIteration: + continue + pager.close() + pager = Pager(source, first) + tab_idx = 0 + scroll = 0 + elif key == ord("\t"): + tab_idx += 1 + scroll = 0 + elif key == curses.KEY_BTAB: + tab_idx -= 1 + scroll = 0 + elif key == curses.KEY_UP: + scroll -= 1 + elif key == curses.KEY_DOWN: + scroll += 1 + elif key == curses.KEY_PPAGE: + scroll -= body_h + elif key == curses.KEY_NPAGE: + scroll += body_h + elif key == curses.KEY_HOME: + scroll = 0 + elif key == curses.KEY_END: + scroll = max_scroll + + +@click.command("inspect") +@click.argument("dataset") +@click.option( + "--split", + default=None, + help="Split to load (auto-detected if omitted).", +) +@click.option( + "--limit", + type=int, + default=None, + help="Load at most this many samples.", +) +@click.option( + "--index", + type=int, + default=0, + show_default=True, + help="Skip this many samples before reading.", +) +def inspect_cmd(dataset, split, limit, index): + """Read a completionist dataset, auto-detecting its format. + + Interactive viewer (or a plain-text dump when stdout is not a TTY) with + per-sample tabs for conversations, prompt/completion, and tools. + """ + if index < 0: + handle_error("--index must be >= 0.") + if limit is not None and limit < 1: + handle_error("--limit must be >= 1.") + + pager = None + try: + tty = sys.stdout.isatty() + if tty: + print("Loading dataset...", file=sys.stderr, flush=True) + start = time.monotonic() + source = DatasetStream(dataset, split, index, limit) + try: + first = source.next() + except StopIteration: + handle_error(f"No samples found in '{dataset}'.") + if tty: + print( + f"done in {time.monotonic() - start:.1f}s", file=sys.stderr, flush=True + ) + + if not sys.stdout.isatty(): + total = source.total + print(render_sample_text(first, detect_format(first), 0, total)) + n = 1 + try: + while True: + row = source.next() + print(render_sample_text(row, detect_format(row), n, total)) + n += 1 + except StopIteration: + pass + return + + pager = Pager(source, first) + curses.wrapper(_tui, source, pager) + except KeyboardInterrupt: + if pager is not None: + pager.close() + # datasets/pyarrow leave background threads that abort the interpreter + # during finalization on Ctrl-C; exit hard to skip the fatal error. + os._exit(130) + finally: + if pager is not None: + pager.close() diff --git a/completionist/commands/topics.py b/completionist/commands/topics.py new file mode 100644 index 0000000..aa40b6f --- /dev/null +++ b/completionist/commands/topics.py @@ -0,0 +1,301 @@ +import json +import os +import sys +import traceback + +import click +from huggingface_hub import get_token +from pydantic import BaseModel, Field + +from completionist.llm_api import get_completion +from completionist.utils import read_file_content + + +class TopicList(BaseModel): + topics: list[str] = Field(..., description="List of conversation topics") + + +DEFAULT_SYSTEM_PROMPT = ( + "You are a creative topic curator. " + "Generate {batch_size} diverse conversation topics across these categories: " + "{categories}. " + "Each topic should be a short phrase that could seed an engaging multi-turn conversation. " + "Avoid generic topics — make them specific and thought-provoking. " + "Return the result as a JSON object matching the schema." +) + + +def _load_existing(path: str, ext: str) -> set[str]: + """Load existing topics from a .txt or .jsonl file.""" + topics: set[str] = set() + if not os.path.exists(path): + return topics + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + if ext == ".jsonl": + try: + data = json.loads(line) + topic = data.get("topic", "") + except json.JSONDecodeError: + continue + else: + topic = line + if topic: + topics.add(topic) + return topics + + +def _write_topics_file(path: str, ext: str, topics: list[str]) -> None: + """Write topics to a .txt (one per line) or .jsonl file.""" + if ext == ".jsonl": + with open(path, "w") as f: + for topic in topics: + f.write(json.dumps({"topic": topic}) + "\n") + else: + with open(path, "w") as f: + for topic in topics: + f.write(topic + "\n") + + +def _generate_batch( + model_name: str, + api_url: str, + system_prompt: str, + batch_size: int, + temperature: float, + top_p: float, + max_tokens: int, + reasoning: str | None, + hf_api_token: str | None, + openai_api_token: str | None, +) -> list[str]: + """Make a single API call and return generated topics.""" + result = get_completion( + prompt=f"Generate {batch_size} topics.", + model_name=model_name, + api_url=api_url, + system_prompt=system_prompt, + hf_api_token=hf_api_token, + openai_api_token=openai_api_token, + pydantic_schema=TopicList, + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + reasoning=reasoning, + ) + + if result is None: + return [] + if isinstance(result, TopicList): + return result.topics + # outlines returned a raw string/dict — parse it as fallback + raw = json.loads(result) if isinstance(result, str) else result + return TopicList(**raw).topics + + +@click.command("topics") +@click.option( + "--output", + "-o", + type=click.Path(dir_okay=False, resolve_path=True), + required=True, + help="Output file path. Extension determines format: .txt (one per line) or .jsonl.", +) +@click.option( + "--num-topics", + "-n", + type=int, + required=True, + help="Total number of unique topics to generate.", +) +@click.option( + "--model-name", + required=True, + help="The name of the model to use for generation.", +) +@click.option( + "--categories", + required=True, + help="Comma-separated list of topic categories to guide generation.", +) +@click.option( + "--api-url", + default="http://localhost:11434/v1", + help="The API endpoint URL for the LLM.", + show_default=True, +) +@click.option( + "--batch-size", + type=int, + default=50, + help="Topics to request per API call. Multiple batches are deduplicated automatically.", + show_default=True, +) +@click.option( + "--temperature", + type=float, + default=0.9, + help="Sampling temperature. Higher values produce more diverse topics.", + show_default=True, +) +@click.option( + "--top-p", + type=float, + default=0.95, + help="Nucleus sampling (top-p).", + show_default=True, +) +@click.option( + "--max-tokens", + type=int, + default=4096, + help="Maximum tokens per API call. Increase for reasoning models.", + show_default=True, +) +@click.option( + "--system-prompt", + default=None, + help="Override the built-in system prompt. Use {batch_size} and {categories} placeholders.", +) +@click.option( + "--system-prompt-file", + type=click.Path(exists=True, dir_okay=False, resolve_path=True), + default=None, + help="Path to a file overriding the built-in system prompt.", +) +@click.option( + "--reasoning", + type=str, + default=None, + help="llama.cpp reasoning mode: 'on', 'off', or 'auto'. Leave unset for model default.", +) +@click.option( + "--max-batches", + type=int, + default=20, + help="Maximum number of API calls before stopping, even if --num-topics is not reached.", + show_default=True, +) +@click.option( + "--append", + is_flag=True, + help="Append to existing output file. --num-topics specifies how many NEW topics to add.", +) +def topics_cmd( + output, + num_topics, + model_name, + categories, + api_url, + batch_size, + temperature, + top_p, + max_tokens, + system_prompt, + system_prompt_file, + reasoning, + max_batches, + append, +): + """Generate diverse conversation topics via LLM. + + Output format is determined by file extension: + \b + .txt — one topic per line + .jsonl — {"topic": "..."} per line + """ + if num_topics < 1: + print("Error: --num-topics must be at least 1.") + sys.exit(1) + + if batch_size < 1: + print("Error: --batch-size must be at least 1.") + sys.exit(1) + + # Resolve system prompt: file > inline > built-in default + system_prompt_content = read_file_content(system_prompt_file) or system_prompt + if not system_prompt_content: + system_prompt_content = DEFAULT_SYSTEM_PROMPT + + system_prompt_content = system_prompt_content.replace( + "{batch_size}", str(batch_size) + ) + system_prompt_content = system_prompt_content.replace("{categories}", categories) + + # Determine output format from extension + ext = os.path.splitext(output)[1].lower() + if ext not in (".txt", ".jsonl"): + print(f"Error: Unsupported output extension '{ext}'. Use .txt or .jsonl.") + sys.exit(1) + + hf_api_token = get_token() + openai_api_token = os.environ.get("OPENAI_API_TOKEN", None) + + # Seed with existing topics when appending + existing_topics: set[str] = set() + all_topics: set[str] = set() + if append and os.path.exists(output): + existing_topics = _load_existing(output, ext) + all_topics.update(existing_topics) + print( + f"Appending to {output} ({len(existing_topics)} existing, adding up to {num_topics} new)" + ) + + batches_run = 0 + target = len(existing_topics) + num_topics # existing + desired new + + def current_topics_list() -> list[str]: + if append: + new_topics = list(all_topics - existing_topics)[:num_topics] + return list(existing_topics) + new_topics + return list(all_topics)[:num_topics] + + print( + f"Generating {num_topics} topics across categories: {categories}\n" + f"Model: {model_name} | API: {api_url} | Batch size: {batch_size}\n" + ) + + while len(all_topics) < target and batches_run < max_batches: + batches_run += 1 + + try: + batch = _generate_batch( + model_name=model_name, + api_url=api_url, + system_prompt=system_prompt_content, + batch_size=batch_size, + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + reasoning=reasoning, + hf_api_token=hf_api_token, + openai_api_token=openai_api_token, + ) + except Exception: + print(f"Batch {batches_run} failed: {traceback.format_exc()}") + continue + + new = set(batch) - all_topics + all_topics.update(batch) + + print( + f" Batch {batches_run}: {len(batch)} generated, " + f"{len(new)} new unique, {len(all_topics)} total" + ) + # Persist progress so an interrupted run can resume via --append. + _write_topics_file(output, ext, current_topics_list()) + + if len(all_topics) < target: + shortfall = target - len(all_topics) + print( + f"\nWarning: reached --max-batches ({max_batches}) with " + f"{len(all_topics)} unique topics — {shortfall} short of {target}." + ) + + topics_list = current_topics_list() + _write_topics_file(output, ext, topics_list) + + print(f"\nSaved {len(topics_list)} unique topics to {output}") diff --git a/completionist/dataset_stream.py b/completionist/dataset_stream.py new file mode 100644 index 0000000..ee04ae1 --- /dev/null +++ b/completionist/dataset_stream.py @@ -0,0 +1,135 @@ +"""Reusable lazy streaming access to local and remote datasets.""" + +import itertools +import os + +from datasets import load_dataset + +from completionist.utils import handle_error + + +def _load(ref): + """Return a streaming dataset handle for a local path or a HF name.""" + if os.path.exists(ref): + if os.path.isdir(ref): + return load_dataset(ref, streaming=True) + ext = os.path.splitext(ref)[1].lower() + if ext == ".parquet": + return load_dataset("parquet", data_files=ref, streaming=True) + if ext in (".jsonl", ".json"): + return load_dataset("json", data_files=ref, streaming=True) + handle_error(f"Unsupported format '{ext}'. Use .parquet or .jsonl.") + return load_dataset(ref, streaming=True) + + +def _pick_split(split, splits): + """Resolve the split to use, honoring an explicit value when given.""" + if split is not None: + if splits and split not in splits: + handle_error(f"Split '{split}' not found. Available: {', '.join(splits)}") + return split + for candidate in ("train", "test", "validation"): + if candidate in splits: + return candidate + return splits[0] if splits else None + + +def _num_examples(ds, split): + """Return the metadata sample count for a split, or None if unknown.""" + if split is None: + return None + info = getattr(ds, "info", None) + if info is None or not info.splits: + return None + split_info = info.splits.get(split) + return split_info.num_examples if split_info else None + + +def _file_row_count(ref): + """Row count from a local parquet footer (metadata only), or None. + + Reads only the parquet footer, never the data, so this does not + double-load the file. jsonl is skipped: counting lines would require + reading the whole file again. + """ + if os.path.splitext(ref)[1].lower() == ".parquet": + import pyarrow.parquet as pq + + return pq.ParquetFile(ref).metadata.num_rows + return None + + +def _trim_total(total, index, limit): + """Adjust a total count for --index and --limit.""" + if total is None: + return None + total = max(0, total - index) + if limit is not None: + total = min(total, limit) + return total + + +class DatasetStream: + """Lazy, switchable streaming source for a local or remote dataset. + + Handles remote HF names, local directories following the HF split layout, + and single local parquet/jsonl files uniformly. Iteration is forward-only; + ``switch_split()`` re-opens the stream for the next split. + """ + + def __init__(self, ref, split=None, index=0, limit=None): + self.ref = ref + self.index = index + self.limit = limit + self._requested = split + self._loaded = _load(ref) + if isinstance(self._loaded, dict): + self.splits = list(self._loaded) + else: + info = getattr(self._loaded, "info", None) + self.splits = list(info.splits) if info is not None and info.splits else [] + self.split = None + self.total = None + self.exhausted = False + self._iterator = None + self._open(_pick_split(self._requested, self.splits)) + + def _open(self, split): + self.split = split + if isinstance(self._loaded, dict): + if split is None: + self._iterator = iter(()) + self.total = 0 + self.exhausted = True + return + ds = self._loaded[split] + else: + ds = self._loaded + stop = None if self.limit is None else self.index + self.limit + self._iterator = itertools.islice(iter(ds), self.index, stop) + total = _num_examples(ds, split) + if total is None and os.path.isfile(self.ref): + total = _file_row_count(self.ref) + self.total = _trim_total(total, self.index, self.limit) + self.exhausted = False + + @property + def cache_key(self): + """Stable key identifying this dataset/split/view, for cache naming.""" + return f"{self.ref}|{self.split}|{self.index}|{self.limit}" + + def next(self): + """Return the next record, raising StopIteration when exhausted.""" + try: + return next(self._iterator) + except StopIteration: + self.exhausted = True + raise + + def switch_split(self): + """Move to the next split, returning False when there is only one.""" + if len(self.splits) < 2: + return False + nxt = self.splits[(self.splits.index(self.split) + 1) % len(self.splits)] + self._open(nxt) + return True diff --git a/completionist/dataset_view.py b/completionist/dataset_view.py new file mode 100644 index 0000000..d74ae9c --- /dev/null +++ b/completionist/dataset_view.py @@ -0,0 +1,136 @@ +"""Dataset format detection and rendering to styled segments (reusable).""" + +import json + +ROLE_PREFIX = {"user": "user: ", "assistant": "assistant: ", "system": "system: "} + + +def detect_format(row): + """Return 'tools', 'conversation', 'prompt_completion', or 'generic'.""" + messages = row.get("messages") + is_messages = ( + isinstance(messages, list) + and bool(messages) + and all(isinstance(m, dict) and "role" in m for m in messages) + ) + has_tools = isinstance(row.get("tools"), list) and bool(row.get("tools")) + has_tool_calls = is_messages and any(m.get("tool_calls") for m in messages) + + if is_messages and (has_tools or has_tool_calls): + return "tools" + if is_messages: + return "conversation" + if "prompt" in row and "completion" in row: + return "prompt_completion" + return "generic" + + +def _pretty_json(value): + return json.dumps(value, indent=2, ensure_ascii=False) + + +def _pretty_args(arguments): + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except (ValueError, TypeError): + return arguments + return json.dumps(arguments, ensure_ascii=False) + + +def _role_style(role): + if role == "assistant": + return "assistant" + if role == "system": + return "system" + return "user" + + +def _conversation_segments(row): + """Turn a conversation row into (style, text) segments.""" + segments = [] + topic = row.get("topic") + if topic: + segments.append(("topic", str(topic))) + segments.append(("plain", "")) + for msg in row.get("messages") or []: + role = msg.get("role", "") + content = msg.get("content") or "" + if role == "tool": + if content: + segments.append( + ("tool", f"tool({msg.get('tool_call_id', '?')}): {content}") + ) + else: + if content: + segments.append((_role_style(role), content)) + for call in msg.get("tool_calls") or []: + fn = call.get("function") or {} + name = fn.get("name", "?") + args = fn.get("arguments") + if args: + segments.append(("tool", f"-> {name}({_pretty_args(args)})")) + else: + segments.append(("tool", f"-> {name}()")) + segments.append(("plain", "")) + return segments + + +def _tools_segments(row): + """Turn the 'tools' field into (style, text) segments.""" + segments = [] + for tool in row.get("tools") or []: + fn = tool.get("function", tool) if isinstance(tool, dict) else {} + segments.append(("topic", fn.get("name", "?"))) + desc = fn.get("description") + if desc: + segments.append(("plain", desc)) + params = fn.get("parameters") + if params is not None: + segments.append(("raw", _pretty_json(params))) + segments.append(("plain", "")) + return segments + + +def row_to_tabs(row, fmt): + """Return ``[(tab_name, [(style, text), ...]), ...]`` for a single row.""" + tabs = [] + if fmt == "conversation": + tabs.append(("conversation", _conversation_segments(row))) + elif fmt == "prompt_completion": + tabs.append(("prompt", [("user", row.get("prompt") or "")])) + tabs.append(("completion", [("assistant", row.get("completion") or "")])) + if "reasoning" in row: + tabs.append(("reasoning", [("topic", row.get("reasoning") or "")])) + elif fmt == "tools": + tabs.append(("conversation", _conversation_segments(row))) + tabs.append(("tools", _tools_segments(row))) + tabs.append(("raw", [("raw", _pretty_json(row))])) + return tabs + + +def title(row, fmt): + """Short one-line label for a row.""" + if fmt in ("conversation", "tools"): + if row.get("topic"): + return str(row["topic"]) + messages = row.get("messages") or [] + if messages: + return str(messages[0].get("content") or "")[:60] + return "" + if fmt == "prompt_completion": + return str(row.get("prompt") or "")[:60] + return "" + + +def render_sample_text(row, fmt, idx, total): + """Render one sample as plain text (used when stdout is not a TTY).""" + header = f"[{idx + 1}/{total if total is not None else '?'}] {title(row, fmt)}" + parts = [header] + for tab_name, segments in row_to_tabs(row, fmt): + parts.append("") + parts.append(f"== {tab_name} ==") + for style, text in segments: + prefix = ROLE_PREFIX.get(style, "") + parts.extend(prefix + ln for ln in text.splitlines() or [""]) + return "\n".join(parts) diff --git a/completionist/pager.py b/completionist/pager.py new file mode 100644 index 0000000..8f61614 --- /dev/null +++ b/completionist/pager.py @@ -0,0 +1,107 @@ +"""Bounded, disk-spilled read-ahead window over a forward-only source.""" + +import collections +import hashlib +import json +import os +import tempfile + +DEFAULT_BUFFER = 200 +DEFAULT_REFILL = 50 + + +class _SpillFile: + """Append-only spill file in /tmp, named by a hash of the source. + + Records evicted from the in-memory window are written here; the file is + only read back when navigating backwards past the window. + """ + + def __init__(self, key): + digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:16] + self.path = os.path.join( + tempfile.gettempdir(), f"completionist-inspect-{digest}.spill" + ) + self._f = open(self.path, "w+b") + + def append(self, row): + offset = self._f.tell() + self._f.write(json.dumps(row, ensure_ascii=False).encode("utf-8") + b"\n") + return offset + + def read(self, offset): + self._f.flush() + self._f.seek(offset) + return json.loads(self._f.readline().decode("utf-8")) + + def remove(self): + self._f.close() + try: + os.unlink(self.path) + except FileNotFoundError: + pass + + +class Pager: + """Bounded sliding window over a forward-only source, disk-spilled. + + Recent records stay in memory; older ones spill to disk so backward + navigation works without holding the whole dataset in RAM. Reads ahead in + chunks (``refill``) only when the cursor nears the window's right edge. + """ + + def __init__( + self, source, first, buffer_size=DEFAULT_BUFFER, refill=DEFAULT_REFILL + ): + self.source = source + self.buffer_size = buffer_size + self.refill = refill + self.buffer = collections.deque([first]) + self.spill = _SpillFile(source.cache_key) + self.offsets = [] # absolute index -> byte offset in the spill file + self.consumed = 1 # records fetched so far + self.cursor = 0 # absolute index of the current row + + def _fetch_one(self): + try: + row = self.source.next() + except StopIteration: + return False + if len(self.buffer) == self.buffer_size: + self.offsets.append(self.spill.append(self.buffer.popleft())) + self.buffer.append(row) + self.consumed += 1 + return True + + def _read(self, i): + mem_base = self.consumed - len(self.buffer) + if i >= mem_base: + return self.buffer[i - mem_base] + return self.spill.read(self.offsets[i]) + + def current(self): + return self._read(self.cursor) + + def next(self): + if self.cursor < self.consumed - 1: + self.cursor += 1 + elif self._fetch_one(): + self.cursor += 1 + if self.consumed - 1 - self.cursor < self.refill: + self._refill() + + def prev(self): + if self.cursor > 0: + self.cursor -= 1 + + def _refill(self): + for _ in range(self.refill): + if not self._fetch_one(): + break + + @property + def total(self): + return self.consumed if self.source.exhausted else self.source.total + + def close(self): + self.spill.remove() diff --git a/completionist/processing.py b/completionist/processing.py index a953ab6..0224214 100644 --- a/completionist/processing.py +++ b/completionist/processing.py @@ -14,6 +14,9 @@ def process_samples_with_executor( """ Manages the concurrent execution of tasks using a ThreadPoolExecutor. + Results are returned in submission order, so callers can resume by index + regardless of worker count. + If save_callback is provided, it is called every save_every completions with the current completions list so partial results can be persisted. @@ -36,7 +39,7 @@ def process_samples_with_executor( futures.append(future) for future in tqdm( - concurrent.futures.as_completed(futures), + futures, total=len(futures), initial=resume_idx, desc="Generating completions", diff --git a/pyproject.toml b/pyproject.toml index 928e284..0f46bd7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "completionist" -version = "0.1.1" +version = "0.1.2" description = "CLI tool for batch LLM operations on datasets — translate, complete, and build structured synthetic datasets" readme = "README.md" requires-python = ">=3.11.11" @@ -54,6 +54,12 @@ convention = "google" # Accepts: "google", "numpy", or "pep257". requires = ["setuptools>=61.0", "wheel", "build"] build-backend = "setuptools.build_meta" +[dependency-groups] +dev = [ + "pre-commit>=4.6.2", + "pytest>=9.1.1", +] + [tool.setuptools.packages.find] where = ["."] include = ["completionist"] diff --git a/tests/test_chat.py b/tests/test_chat.py new file mode 100644 index 0000000..0200b04 --- /dev/null +++ b/tests/test_chat.py @@ -0,0 +1,225 @@ +import time + +from click.testing import CliRunner + +from completionist.commands import chat +from completionist.commands.chat import ChatConversation, ChatMessage + + +def _llm_config(**overrides): + config = { + "model_name": "test-model", + "api_url": "http://localhost:11434/v1", + "system_prompt": "exactly {num_turns} messages", + "user_prompt_template": "topic: {topic} with {num_turns} turns", + "min_turns": 4, + "max_turns": 6, + "generation_config": {"temperature": 0.7, "top_p": 0.95}, + "hf_api_token": None, + "openai_api_token": None, + "reasoning": None, + "max_tokens": 4096, + } + config.update(overrides) + return config + + +def _conversation(): + return ChatConversation( + topic="test topic", + messages=[ + ChatMessage(role="user", content="hi"), + ChatMessage(role="assistant", content="hello"), + ], + ) + + +def test_chat_conversation_schema(): + conv = _conversation() + assert conv.topic == "test topic" + assert conv.messages[0].role == "user" + + +def test_task_handler_uses_even_turn_count(monkeypatch): + captured = {} + + def fake_get_completion(**kwargs): + captured.update(kwargs) + return _conversation() + + monkeypatch.setattr(chat, "get_completion", fake_get_completion) + result = chat.chat_task_handler("my topic", _llm_config()) + + system_prompt = captured["system_prompt"] + assert ( + "exactly 4 messages" in system_prompt or "exactly 6 messages" in system_prompt + ) + assert "exactly 5 messages" not in system_prompt + assert result["topic"] == "test topic" + + +def test_task_handler_returns_none_on_failure(monkeypatch): + monkeypatch.setattr(chat, "get_completion", lambda **kwargs: None) + result = chat.chat_task_handler("my topic", _llm_config()) + assert result is None + + +def test_task_handler_ignores_unknown_placeholder(monkeypatch): + # A custom prompt with a placeholder the handler doesn't know must not crash. + monkeypatch.setattr(chat, "get_completion", lambda **kwargs: _conversation()) + result = chat.chat_task_handler( + "my topic", + _llm_config( + system_prompt="literal {unknown} and {num_turns}", + user_prompt_template="{topic} and {unknown}", + ), + ) + assert result is not None + + +def test_chat_cmd_resumes_from_existing_output(tmp_path, monkeypatch): + topics_file = tmp_path / "topics.txt" + topics_file.write_text("topic A\n") + output_file = tmp_path / "out.parquet" + output_file.touch() + + existing = [ + { + "topic": "topic A", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hey"}, + ], + }, + { + "topic": "topic A", + "messages": [ + {"role": "user", "content": "hi2"}, + {"role": "assistant", "content": "hey2"}, + ], + }, + ] + + class FakeDataset: + def __init__(self): + self.rows = existing + + @classmethod + def from_parquet(cls, path): + return cls() + + @classmethod + def from_json(cls, path): + return cls() + + def to_list(self): + return self.rows + + monkeypatch.setattr(chat, "Dataset", FakeDataset) + monkeypatch.setattr(chat, "get_token", lambda: None) + + captured = {} + + def fake_process( + dataset_to_process, + workers, + resume_idx, + task_handler, + llm_config, + save_callback=None, + save_every=50, + ): + captured["dataset"] = list(dataset_to_process) + captured["resume_idx"] = resume_idx + captured["save_callback"] = save_callback + return [ + { + "topic": "topic A", + "messages": [ + {"role": "user", "content": "new"}, + {"role": "assistant", "content": "new2"}, + ], + } + ] + + monkeypatch.setattr(chat, "process_samples_with_executor", fake_process) + + saved = {} + + def fake_save(completions, output_file, push_to_hub, hf_repo_id, hf_api_token): + saved["completions"] = completions + + monkeypatch.setattr(chat, "save_and_push_dataset", fake_save) + + runner = CliRunner() + result = runner.invoke( + chat.chat_cmd, + [ + "--topics-file", + str(topics_file), + "--num-conversations", + "3", + "--output-file", + str(output_file), + "--model-name", + "test-model", + "--workers", + "1", + ], + ) + + assert result.exit_code == 0, result.output + # tasks = [topic A] * 3; 2 already done -> 1 remaining task. + assert captured["resume_idx"] == 2 + assert captured["dataset"] == ["topic A"] + # final save merges existing (2) + newly generated (1). + assert len(saved["completions"]) == 3 + + +def test_task_handler_retries_transient_failure(monkeypatch): + calls = {"n": 0} + + def flaky_get_completion(**kwargs): + calls["n"] += 1 + if calls["n"] < 3: + return None + return _conversation() + + monkeypatch.setattr(chat, "get_completion", flaky_get_completion) + monkeypatch.setattr(time, "sleep", lambda _: None) + + result = chat.chat_task_handler("my topic", _llm_config(retries=3)) + assert result is not None + assert calls["n"] == 3 + + +def test_task_handler_gives_up_after_retries(monkeypatch): + calls = {"n": 0} + + def always_none(**kwargs): + calls["n"] += 1 + return None + + monkeypatch.setattr(chat, "get_completion", always_none) + monkeypatch.setattr(time, "sleep", lambda _: None) + + result = chat.chat_task_handler("my topic", _llm_config(retries=2)) + assert result is None + assert calls["n"] == 2 + + +def test_task_handler_retries_invalid_json(monkeypatch): + calls = {"n": 0} + + def flaky(**kwargs): + calls["n"] += 1 + if calls["n"] == 1: + return "this is not json" + return _conversation() + + monkeypatch.setattr(chat, "get_completion", flaky) + monkeypatch.setattr(time, "sleep", lambda _: None) + + result = chat.chat_task_handler("my topic", _llm_config(retries=3)) + assert result is not None + assert calls["n"] == 2 diff --git a/tests/test_dataset_stream.py b/tests/test_dataset_stream.py new file mode 100644 index 0000000..257a0d5 --- /dev/null +++ b/tests/test_dataset_stream.py @@ -0,0 +1,192 @@ +import pytest +from datasets import Dataset + +from completionist import dataset_stream +from completionist.dataset_stream import DatasetStream + + +def _fake_dataset(split_data): + """Build a fake streaming load_dataset result: {split: iterable-with-info}.""" + + class _SplitInfo: + def __init__(self, n): + self.num_examples = n + + class _Info: + def __init__(self, name, n): + self.splits = {name: _SplitInfo(n)} + + class _Iter: + def __init__(self, rows, info): + self._rows = rows + self.info = info + + def __iter__(self): + return iter(self._rows) + + return {k: _Iter(v, _Info(k, len(v))) for k, v in split_data.items()} + + +def test_load_local_parquet(tmp_path): + path = tmp_path / "data.parquet" + Dataset.from_list( + [{"prompt": str(i), "completion": str(i)} for i in range(5)] + ).to_parquet(str(path)) + src = DatasetStream(str(path)) + assert src.splits == ["train"] + # Total comes from the parquet footer (metadata only, no data re-read). + assert src.total == 5 + rows = [] + while True: + try: + rows.append(src.next()) + except StopIteration: + break + assert len(rows) == 5 + assert src.exhausted + + +def test_load_local_jsonl(tmp_path): + path = tmp_path / "data.jsonl" + Dataset.from_list( + [{"prompt": str(i), "completion": str(i)} for i in range(3)] + ).to_json(str(path)) + src = DatasetStream(str(path)) + assert src.total is None + assert src.next() == {"prompt": "0", "completion": "0"} + + +def test_load_local_directory_with_splits(tmp_path): + (tmp_path / "train").mkdir() + (tmp_path / "test").mkdir() + Dataset.from_list([{"x": i} for i in range(4)]).to_parquet( + str(tmp_path / "train" / "part-00000.parquet") + ) + Dataset.from_list([{"x": i} for i in range(2)]).to_parquet( + str(tmp_path / "test" / "part-00000.parquet") + ) + src = DatasetStream(str(tmp_path)) + assert src.splits == ["train", "test"] + assert src.split == "train" + assert src.total is None + assert src.switch_split() is True + assert src.split == "test" + assert src.total is None + + +def test_remote_splits_and_total(monkeypatch): + monkeypatch.setattr( + dataset_stream, + "load_dataset", + lambda *a, **k: _fake_dataset( + {"train": [{"x": 1}, {"x": 2}], "test": [{"x": 3}]} + ), + ) + src = DatasetStream("some/dataset") + assert src.splits == ["train", "test"] + assert src.split == "train" + assert src.total == 2 + assert src.next() == {"x": 1} + assert src.next() == {"x": 2} + assert src.switch_split() is True + assert src.split == "test" + assert src.total == 1 + assert src.next() == {"x": 3} + + +def test_remote_single_split(monkeypatch): + monkeypatch.setattr( + dataset_stream, + "load_dataset", + lambda *a, **k: _fake_dataset({"train": [{"x": 1}]}), + ) + src = DatasetStream("some/dataset") + assert src.switch_split() is False + + +def test_remote_index_and_limit(monkeypatch): + rows = [{"x": i} for i in range(10)] + monkeypatch.setattr( + dataset_stream, "load_dataset", lambda *a, **k: _fake_dataset({"train": rows}) + ) + src = DatasetStream("some/dataset", index=3, limit=4) + assert src.total == 4 + got = [] + while True: + try: + got.append(src.next()) + except StopIteration: + break + assert got == [{"x": 3}, {"x": 4}, {"x": 5}, {"x": 6}] + + +def test_pick_split_explicit(): + assert dataset_stream._pick_split("test", ["train", "test"]) == "test" + + +def test_pick_split_auto_prefers_train(): + assert dataset_stream._pick_split(None, ["test", "train"]) == "train" + + +def test_pick_split_auto_falls_back(): + assert dataset_stream._pick_split(None, ["validation"]) == "validation" + + +def test_pick_split_auto_empty(): + assert dataset_stream._pick_split(None, []) is None + + +def test_pick_split_explicit_missing_errors(): + with pytest.raises(SystemExit): + dataset_stream._pick_split("nope", ["train"]) + + +def test_trim_total(): + assert dataset_stream._trim_total(10, 2, 3) == 3 + assert dataset_stream._trim_total(10, 2, None) == 8 + assert dataset_stream._trim_total(10, 20, None) == 0 + assert dataset_stream._trim_total(None, 0, None) is None + + +def test_unsupported_extension_errors(tmp_path): + path = tmp_path / "data.txt" + path.write_text("hello") + with pytest.raises(SystemExit): + DatasetStream(str(path)) + + +def test_cache_key_changes_with_split(monkeypatch): + monkeypatch.setattr( + dataset_stream, + "load_dataset", + lambda *a, **k: _fake_dataset({"train": [{"x": 1}], "test": [{"x": 2}]}), + ) + src = DatasetStream("some/dataset", index=3, limit=4) + key = src.cache_key + assert "some/dataset" in key + assert "train" in key + src.switch_split() + assert src.cache_key != key + assert "test" in src.cache_key + + +def test_num_examples_returns_none_without_info(): + class _NoInfo: + pass + + assert dataset_stream._num_examples(_NoInfo(), "train") is None + assert dataset_stream._num_examples(_NoInfo(), None) is None + + +def test_switch_split_cycles(monkeypatch): + monkeypatch.setattr( + dataset_stream, + "load_dataset", + lambda *a, **k: _fake_dataset({"a": [1], "b": [2]}), + ) + src = DatasetStream("some/dataset") + assert src.split == "a" + assert src.switch_split() is True + assert src.split == "b" + assert src.switch_split() is True + assert src.split == "a" diff --git a/tests/test_dataset_view.py b/tests/test_dataset_view.py new file mode 100644 index 0000000..e988bed --- /dev/null +++ b/tests/test_dataset_view.py @@ -0,0 +1,152 @@ +from completionist.dataset_view import ( + detect_format, + render_sample_text, + row_to_tabs, + title, +) + + +def _conversation_row(): + return { + "topic": "climate policy", + "messages": [ + {"role": "user", "content": "Is degrowth viable?"}, + {"role": "assistant", "content": "Partly, but it needs a just transition."}, + ], + } + + +def _prompt_completion_row(): + return {"prompt": "What is justice?", "completion": "A contested concept."} + + +def _tools_row(): + return { + "topic": "weather", + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } + ], + "messages": [ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "SF"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": '{"temp": 20}'}, + ], + } + + +def test_detect_format_conversation(): + assert detect_format(_conversation_row()) == "conversation" + + +def test_detect_format_prompt_completion(): + assert detect_format(_prompt_completion_row()) == "prompt_completion" + + +def test_detect_format_tools_with_tools_field(): + assert detect_format(_tools_row()) == "tools" + + +def test_detect_format_tools_without_tools_field(): + row = _conversation_row() + row["messages"].append( + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "c", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + } + ) + assert detect_format(row) == "tools" + + +def test_detect_format_generic(): + assert detect_format({"foo": "bar"}) == "generic" + + +def test_row_to_tabs_conversation(): + tabs = row_to_tabs(_conversation_row(), "conversation") + names = [name for name, _ in tabs] + assert names == ["conversation", "raw"] + segments = tabs[0][1] + assert ("user", "Is degrowth viable?") in segments + assert ("assistant", "Partly, but it needs a just transition.") in segments + + +def test_row_to_tabs_prompt_completion_with_reasoning(): + row = _prompt_completion_row() + row["reasoning"] = "thinking..." + tabs = row_to_tabs(row, "prompt_completion") + assert [name for name, _ in tabs] == ["prompt", "completion", "reasoning", "raw"] + + +def test_row_to_tabs_prompt_completion_without_reasoning(): + tabs = row_to_tabs(_prompt_completion_row(), "prompt_completion") + assert [name for name, _ in tabs] == ["prompt", "completion", "raw"] + + +def test_row_to_tabs_tools(): + tabs = row_to_tabs(_tools_row(), "tools") + names = [name for name, _ in tabs] + assert names == ["conversation", "tools", "raw"] + + conv_text = [text for _, text in tabs[0][1]] + assert any(text.startswith("-> get_weather(") for text in conv_text) + assert any("tool(call_1)" in text for text in conv_text) + + tools_text = [text for _, text in tabs[1][1]] + assert "get_weather" in tools_text + assert any('"city"' in text for text in tools_text) + + +def test_title(): + assert title(_conversation_row(), "conversation") == "climate policy" + assert title(_prompt_completion_row(), "prompt_completion") == "What is justice?" + + +def test_render_sample_text(): + out = render_sample_text(_conversation_row(), "conversation", idx=0, total=1) + assert "[1/1] climate policy" in out + assert "== conversation ==" in out + assert "user: Is degrowth viable?" in out + assert "assistant: Partly, but it needs a just transition." in out + assert "== raw ==" in out + + +def test_render_sample_text_unknown_total(): + out = render_sample_text( + _prompt_completion_row(), "prompt_completion", idx=2, total=None + ) + assert "[3/?]" in out + + +def test_row_to_tabs_generic(): + tabs = row_to_tabs({"foo": "bar"}, "generic") + assert [name for name, _ in tabs] == ["raw"] diff --git a/tests/test_inspect.py b/tests/test_inspect.py new file mode 100644 index 0000000..95d97ad --- /dev/null +++ b/tests/test_inspect.py @@ -0,0 +1,157 @@ +from click.testing import CliRunner +from datasets import Dataset + +from completionist import dataset_stream as ds_module +from completionist.commands import inspect + + +def _conversation_row(): + return { + "topic": "climate policy", + "messages": [ + {"role": "user", "content": "Is degrowth viable?"}, + {"role": "assistant", "content": "Partly, but it needs a just transition."}, + ], + } + + +def _prompt_completion_row(): + return {"prompt": "What is justice?", "completion": "A contested concept."} + + +def _fake_dataset(split_data): + """Build a fake streaming load_dataset result: {split: iterable-with-info}.""" + + class _SplitInfo: + def __init__(self, n): + self.num_examples = n + + class _Info: + def __init__(self, name, n): + self.splits = {name: _SplitInfo(n)} + + class _Iter: + def __init__(self, rows, info): + self._rows = rows + self.info = info + + def __iter__(self): + return iter(self._rows) + + return {k: _Iter(v, _Info(k, len(v))) for k, v in split_data.items()} + + +def test_inspect_dumps_conversation(tmp_path): + path = tmp_path / "conv.parquet" + Dataset.from_list([_conversation_row()]).to_parquet(str(path)) + + result = CliRunner().invoke(inspect.inspect_cmd, [str(path)]) + assert result.exit_code == 0, result.output + assert "== conversation ==" in result.output + assert "user: Is degrowth viable?" in result.output + assert "assistant: Partly, but it needs a just transition." in result.output + assert "== raw ==" in result.output + + +def test_inspect_dumps_prompt_completion(tmp_path): + path = tmp_path / "pc.parquet" + Dataset.from_list([_prompt_completion_row()]).to_parquet(str(path)) + + result = CliRunner().invoke(inspect.inspect_cmd, [str(path)]) + assert result.exit_code == 0, result.output + assert "== prompt ==" in result.output + assert "== completion ==" in result.output + assert "user: What is justice?" in result.output + + +def test_inspect_dumps_local_total_from_footer(tmp_path): + # Total for a local parquet comes from its footer metadata (no re-read). + path = tmp_path / "many.parquet" + Dataset.from_list( + [{"prompt": str(i), "completion": str(i)} for i in range(3)] + ).to_parquet(str(path)) + + result = CliRunner().invoke(inspect.inspect_cmd, [str(path)]) + assert result.exit_code == 0, result.output + assert "[1/3]" in result.output + assert "[3/3]" in result.output + + +def test_inspect_dumps_remote_total_from_metadata(monkeypatch): + monkeypatch.setattr( + ds_module, + "load_dataset", + lambda *a, **k: _fake_dataset( + { + "train": [ + {"prompt": "a", "completion": "1"}, + {"prompt": "b", "completion": "2"}, + ] + } + ), + ) + + result = CliRunner().invoke(inspect.inspect_cmd, ["some/dataset"]) + assert result.exit_code == 0, result.output + assert "[1/2]" in result.output + assert "[2/2]" in result.output + + +def test_inspect_dumps_with_index_and_limit(tmp_path): + path = tmp_path / "many.parquet" + Dataset.from_list( + [{"prompt": str(i), "completion": str(i)} for i in range(5)] + ).to_parquet(str(path)) + + result = CliRunner().invoke( + inspect.inspect_cmd, [str(path), "--index", "1", "--limit", "2"] + ) + assert result.exit_code == 0, result.output + assert '"prompt": "1"' in result.output + assert '"prompt": "2"' in result.output + assert '"prompt": "4"' not in result.output + + +def test_inspect_errors_on_unsupported_extension(tmp_path): + path = tmp_path / "data.txt" + path.write_text("hello") + result = CliRunner().invoke(inspect.inspect_cmd, [str(path)]) + assert result.exit_code == 1 + assert "Unsupported format" in result.output + + +def test_inspect_errors_on_empty_dataset(monkeypatch): + monkeypatch.setattr( + ds_module, "load_dataset", lambda *a, **k: _fake_dataset({"train": []}) + ) + result = CliRunner().invoke(inspect.inspect_cmd, ["some/dataset"]) + assert result.exit_code == 1 + assert "No samples found" in result.output + + +def test_inspect_errors_on_invalid_index(tmp_path): + path = tmp_path / "data.parquet" + Dataset.from_list([_prompt_completion_row()]).to_parquet(str(path)) + result = CliRunner().invoke(inspect.inspect_cmd, [str(path), "--index", "-1"]) + assert result.exit_code == 1 + + +def test_wrap_segments_hanging_indent(): + lines = inspect._wrap_segments( + [("user", "a long message that definitely wraps")], width=20 + ) + assert lines[0][0] == "user" + assert lines[0][1].startswith("user: ") + for style, text in lines[1:]: + assert style == "user" + assert text.startswith(" " * len("user: ")) + + +def test_wrap_segments_raw_verbatim(): + lines = inspect._wrap_segments([("raw", "line1\nline2")], width=80) + assert [text for _, text in lines] == ["line1", "line2"] + + +def test_wrap_segments_plain_wraps(): + lines = inspect._wrap_segments([("plain", "aaaa bbbb cccc")], width=5) + assert lines == [("plain", "aaaa"), ("plain", "bbbb"), ("plain", "cccc")] diff --git a/tests/test_pager.py b/tests/test_pager.py new file mode 100644 index 0000000..3979e62 --- /dev/null +++ b/tests/test_pager.py @@ -0,0 +1,91 @@ +import os + +from completionist.pager import Pager + + +class _FakeSource: + def __init__(self, data): + self._it = iter(data) + self.total = len(data) + self.exhausted = False + self.cache_key = "test-key" + + def next(self): + try: + return next(self._it) + except StopIteration: + self.exhausted = True + raise + + +def _pager(data, buffer_size, refill): + """Build a Pager like inspect_cmd does: first record consumed from source.""" + src = _FakeSource(data) + first = src.next() + return Pager(src, first, buffer_size=buffer_size, refill=refill) + + +def test_forward_navigation(): + pager = _pager([{"x": i} for i in range(6)], buffer_size=3, refill=1) + for _ in range(5): + pager.next() + assert pager.current() == {"x": 5} + pager.close() + + +def test_backward_reads_spill(): + pager = _pager([{"x": i} for i in range(6)], buffer_size=3, refill=1) + for _ in range(5): + pager.next() + for i in range(5, -1, -1): + assert pager.current() == {"x": i} + pager.prev() + pager.close() + + +def test_total_on_exhaustion(): + pager = _pager([{"x": i} for i in range(3)], buffer_size=2, refill=2) + prev = -1 + while pager.cursor != prev: + prev = pager.cursor + pager.next() + assert pager.source.exhausted + assert pager.total == 3 + pager.close() + + +def test_close_removes_spill(): + pager = _pager([{"x": i} for i in range(3)], buffer_size=2, refill=2) + path = pager.spill.path + assert os.path.exists(path) + pager.close() + assert not os.path.exists(path) + + +def test_prev_stays_at_start(): + pager = _pager([{"x": i} for i in range(3)], buffer_size=2, refill=2) + pager.prev() + assert pager.cursor == 0 + pager.close() + + +def test_next_stays_at_end(): + pager = _pager([{"x": i} for i in range(3)], buffer_size=2, refill=2) + for _ in range(3): + pager.next() + assert pager.cursor == 2 + pager.next() + assert pager.cursor == 2 + pager.close() + + +def test_buffer_stays_bounded(): + pager = _pager([{"x": i} for i in range(30)], buffer_size=10, refill=5) + prev = -1 + while pager.cursor != prev: + prev = pager.cursor + pager.next() + assert pager.source.exhausted + assert pager.total == 30 + assert len(pager.buffer) <= 10 + pager.close() diff --git a/tests/test_processing.py b/tests/test_processing.py new file mode 100644 index 0000000..104e6e1 --- /dev/null +++ b/tests/test_processing.py @@ -0,0 +1,76 @@ +import random +import time + +from completionist.processing import process_samples_with_executor + + +def _slow_handler(sample, llm_config): + # Later samples finish faster, forcing out-of-order completion. + time.sleep(0.005 * (50 - sample)) + return sample + + +def test_preserves_submission_order(): + data = list(range(1, 21)) + results = process_samples_with_executor( + dataset_to_process=data, + workers=4, + resume_idx=0, + task_handler=_slow_handler, + llm_config={}, + ) + assert results == data + + +def test_multi_worker_complete_and_ordered(): + rng = random.Random(0) + data = list(range(1, 101)) + delays = {x: rng.uniform(0.0, 0.02) for x in data} + + def handler(sample, cfg): + time.sleep(delays[sample]) + return sample + + results = process_samples_with_executor( + dataset_to_process=data, + workers=8, + resume_idx=0, + task_handler=handler, + llm_config={}, + ) + assert results == data + + +def test_save_callback_checkpoints(): + data = list(range(1, 31)) + saved = [] + + def callback(completions): + saved.append(list(completions)) + + process_samples_with_executor( + dataset_to_process=data, + workers=2, + resume_idx=0, + task_handler=lambda x, cfg: x, + llm_config={}, + save_callback=callback, + save_every=10, + ) + assert saved == [list(range(1, 11)), list(range(1, 21)), list(range(1, 31))] + + +def test_none_results_are_skipped(): + data = ["keep", "skip", "keep"] + + def handler(sample, cfg): + return None if sample == "skip" else sample + + results = process_samples_with_executor( + dataset_to_process=data, + workers=1, + resume_idx=0, + task_handler=handler, + llm_config={}, + ) + assert results == ["keep", "keep"] diff --git a/tests/test_topics.py b/tests/test_topics.py new file mode 100644 index 0000000..c13a749 --- /dev/null +++ b/tests/test_topics.py @@ -0,0 +1,71 @@ +from completionist.commands import topics +from completionist.commands.topics import TopicList + + +def _call_generate_batch(monkeypatch, result): + monkeypatch.setattr(topics, "get_completion", lambda **kwargs: result) + return topics._generate_batch( + model_name="test-model", + api_url="http://localhost:11434/v1", + system_prompt="generate {batch_size} topics about {categories}", + batch_size=5, + temperature=0.9, + top_p=0.95, + max_tokens=4096, + reasoning=None, + hf_api_token=None, + openai_api_token=None, + ) + + +def test_load_existing_txt(tmp_path): + f = tmp_path / "topics.txt" + f.write_text("topic one\ntopic two\n\ntopic one\n") + assert topics._load_existing(str(f), ".txt") == {"topic one", "topic two"} + + +def test_load_existing_jsonl(tmp_path): + f = tmp_path / "topics.jsonl" + f.write_text('{"topic": "a"}\n{"topic": "b"}\nnot-json\n') + assert topics._load_existing(str(f), ".jsonl") == {"a", "b"} + + +def test_load_existing_missing_file(tmp_path): + assert topics._load_existing(str(tmp_path / "nope.txt"), ".txt") == set() + + +def test_write_topics_file_txt(tmp_path): + f = tmp_path / "topics.txt" + topics._write_topics_file(str(f), ".txt", ["a", "b", "c"]) + assert f.read_text() == "a\nb\nc\n" + + +def test_write_topics_file_jsonl(tmp_path): + f = tmp_path / "topics.jsonl" + topics._write_topics_file(str(f), ".jsonl", ["a", "b"]) + assert f.read_text() == '{"topic": "a"}\n{"topic": "b"}\n' + + +def test_write_then_load_roundtrip(tmp_path): + f = tmp_path / "topics.txt" + topics._write_topics_file(str(f), ".txt", ["x", "y", "z"]) + assert topics._load_existing(str(f), ".txt") == {"x", "y", "z"} + + +def test_topic_list_schema(): + assert TopicList(topics=["a", "b"]).topics == ["a", "b"] + + +def test_generate_batch_returns_topics(monkeypatch): + result = _call_generate_batch(monkeypatch, TopicList(topics=["a", "b"])) + assert result == ["a", "b"] + + +def test_generate_batch_returns_empty_on_none(monkeypatch): + result = _call_generate_batch(monkeypatch, None) + assert result == [] + + +def test_generate_batch_fallback_json_string(monkeypatch): + result = _call_generate_batch(monkeypatch, '{"topics": ["a", "b"]}') + assert result == ["a", "b"] diff --git a/uv.lock b/uv.lock index c1e9dd1..5f6eb1f 100644 --- a/uv.lock +++ b/uv.lock @@ -137,6 +137,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4f/52/34c6cf5bb9285074dc3531c437b3919e825d976fde097a7a73f79e726d03/certifi-2025.7.14-py3-none-any.whl", hash = "sha256:6b31f564a415d79ee77df69d757bb49a5bb53bd9f756cbbe24394ffd6fc1f4b2", size = 162722, upload-time = "2025-07-14T03:29:26.863Z" }, ] +[[package]] +name = "cfgv" +version = "3.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/b5/721b8799b04bf9afe054a3899c6cf4e880fcf8563cc71c15610242490a0c/cfgv-3.5.0.tar.gz", hash = "sha256:d5b1034354820651caa73ede66a6294d6e95c1b00acc5e9b098e917404669132", size = 7334, upload-time = "2025-11-19T20:55:51.612Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl", hash = "sha256:a8dc6b26ad22ff227d2634a65cb388215ce6cc96bbcc5cfde7641ae87e8dacc0", size = 7445, upload-time = "2025-11-19T20:55:50.744Z" }, +] + [[package]] name = "charset-normalizer" version = "3.4.2" @@ -228,6 +237,12 @@ dependencies = [ { name = "tqdm" }, ] +[package.dev-dependencies] +dev = [ + { name = "pre-commit" }, + { name = "pytest" }, +] + [package.metadata] requires-dist = [ { name = "click", specifier = ">=8.3.0" }, @@ -238,6 +253,12 @@ requires-dist = [ { name = "tqdm", specifier = ">=4.67.1" }, ] +[package.metadata.requires-dev] +dev = [ + { name = "pre-commit", specifier = ">=4.6.2" }, + { name = "pytest", specifier = ">=9.1.1" }, +] + [[package]] name = "datasets" version = "5.0.0" @@ -281,6 +302,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl", hash = "sha256:5e31b2d5fbad117cc363ebaf6b689474db18a1f6438bc82358b024abd4c2ca19", size = 45550, upload-time = "2023-08-31T06:11:58.822Z" }, ] +[[package]] +name = "distlib" +version = "0.4.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c9/02/bd72be9134d25ed783ecbbc38a539ffaefbf90c78418c7fb7229600dbac7/distlib-0.4.3.tar.gz", hash = "sha256:f152097224a0ae24be5a0f6bae1b9359af82133bce63f98a95f86cae1aede9ed", size = 615141, upload-time = "2026-06-12T08:04:52.847Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/08/9c41fb51ab5b43eb21674aff13df270e8ba6c4b29c8624e328dc7a9482af/distlib-0.4.3-py2.py3-none-any.whl", hash = "sha256:4b0ce306c966eb73bc3a7b6abad017c556dadd92c44701562cd528ac7fde4d5b", size = 470628, upload-time = "2026-06-12T08:04:50.506Z" }, +] + [[package]] name = "distro" version = "1.9.0" @@ -292,11 +322,11 @@ wheels = [ [[package]] name = "filelock" -version = "3.18.0" +version = "3.32.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0a/10/c23352565a6544bdc5353e0b15fc1c563352101f30e24bf500207a54df9a/filelock-3.18.0.tar.gz", hash = "sha256:adbc88eabb99d2fec8c9c1b229b171f18afa655400173ddc653d5d01501fb9f2", size = 18075, upload-time = "2025-03-14T07:11:40.47Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6d/30/03b03951873a1a0ffc7e8ca0e10c15597b59e8d0e39260704cd2ea087bc4/filelock-3.32.4.tar.gz", hash = "sha256:2bde2e4cf732e0153406d8a7bc80620ecf5e621fe0d25e41143c4e3b4733ff30", size = 222126, upload-time = "2026-08-23T17:37:55.363Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4d/36/2a115987e2d8c300a974597416d9de88f2444426de9571f4b59b2cca3acc/filelock-3.18.0-py3-none-any.whl", hash = "sha256:c401f4f8377c4464e6db25fff06205fd89bdd83b65eb0488ed1b160f780e21de", size = 16215, upload-time = "2025-03-14T07:11:39.145Z" }, + { url = "https://files.pythonhosted.org/packages/01/a4/9b63d595d748e3aff8812b65eacc1a2c4bd90b7c2012e08e72373b4835eb/filelock-3.32.4-py3-none-any.whl", hash = "sha256:22e58ca3b1ae3b98993b762d7338367ae64fe50252bf78d59da3bfebcdf1cedd", size = 99864, upload-time = "2026-08-23T17:37:53.913Z" }, ] [[package]] @@ -470,6 +500,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/59/a8/4677014e771ed1591a87b63a2392ce6923baf807193deef302dcfde17542/huggingface_hub-0.34.3-py3-none-any.whl", hash = "sha256:5444550099e2d86e68b2898b09e85878fbd788fc2957b506c6a79ce060e39492", size = 558847, upload-time = "2025-07-29T08:38:51.904Z" }, ] +[[package]] +name = "identify" +version = "2.6.19" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/52/63/51723b5f116cc04b061cb6f5a561790abf249d25931d515cd375e063e0f4/identify-2.6.19.tar.gz", hash = "sha256:6be5020c38fcb07da56c53733538a3081ea5aa70d36a156f83044bfbf9173842", size = 99567, upload-time = "2026-04-17T18:39:50.265Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/84/d9273cd09688070a6523c4aee4663a8538721b2b755c4962aafae0011e72/identify-2.6.19-py2.py3-none-any.whl", hash = "sha256:20e6a87f786f768c092a721ad107fc9df0eb89347be9396cadf3f4abbd1fb78a", size = 99397, upload-time = "2026-04-17T18:39:49.221Z" }, +] + [[package]] name = "idna" version = "3.10" @@ -479,6 +518,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442, upload-time = "2024-09-15T18:07:37.964Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -735,6 +783,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/d9/f7f9379981e39b8c2511c9e0326d212accacb82f12fbfdc1aa2ce2a7b2b6/multiprocess-0.70.16-py39-none-any.whl", hash = "sha256:a0bafd3ae1b732eac64be2e72038231c1ba97724b60b09400d68f229fcc2fbf3", size = 133351, upload-time = "2024-01-28T18:52:31.981Z" }, ] +[[package]] +name = "nodeenv" +version = "1.10.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/24/bf/d1bda4f6168e0b2e9e5958945e01910052158313224ada5ce1fb2e1113b8/nodeenv-1.10.0.tar.gz", hash = "sha256:996c191ad80897d076bdfba80a41994c2b47c68e224c542b48feba42ba00f8bb", size = 55611, upload-time = "2025-12-20T14:08:54.006Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" }, +] + [[package]] name = "numpy" version = "2.3.2" @@ -1030,6 +1087,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/34/e7/ae39f538fd6844e982063c3a5e4598b8ced43b9633baa3a85ef33af8c05c/pillow-11.3.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c84d689db21a1c397d001aa08241044aa2069e7587b398c8cc63020390b1c1b8", size = 6984598, upload-time = "2025-07-01T09:16:27.732Z" }, ] +[[package]] +name = "platformdirs" +version = "4.11.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/50/bb/ebc6636e1ae41314f796ebb7215fd28febb45f9aac72f2b04cb74b5071dc/platformdirs-4.11.4.tar.gz", hash = "sha256:f3373be828247211d0febabea97e238c3dfde8a60b3c90c32756fb52cb21556d", size = 34079, upload-time = "2026-08-24T14:53:49.676Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/28/be/0ff05fcd2938fb58ad9219bd54135968342d214737e012d62d43f06a2dd6/platformdirs-4.11.4-py3-none-any.whl", hash = "sha256:e34ff91a24bcddc6d939b878bdf3f5c437c9c46fe9e212b1bf455fdf1ee57586", size = 23741, upload-time = "2026-08-24T14:53:48.406Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "ply" version = "3.11" @@ -1039,6 +1114,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a3/58/35da89ee790598a0700ea49b2a66594140f44dec458c07e8e3d4979137fc/ply-3.11-py2.py3-none-any.whl", hash = "sha256:096f9b8350b65ebd2fd1346b12452efe5b9607f7482813ffca50c22722a807ce", size = 49567, upload-time = "2018-02-15T19:01:27.172Z" }, ] +[[package]] +name = "pre-commit" +version = "4.6.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cfgv" }, + { name = "identify" }, + { name = "nodeenv" }, + { name = "pyyaml" }, + { name = "virtualenv" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/74/89/1f3e8e1fc3e97de0fa963495832f581f025f29471602a309e48808244292/pre_commit-4.6.2.tar.gz", hash = "sha256:8f5d7bfb021ecdbcd9d49d89847082dd24172ccde534390081a679ad046e2441", size = 198670, upload-time = "2026-08-10T22:07:18.421Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/e2/bbb7129c9e7999a6b8ee9cca3b66486c25c423ab5a75f34071798b74ce94/pre_commit-4.6.2-py2.py3-none-any.whl", hash = "sha256:e2dde9a75d3bce11bd3831c26d134df00a2803c1d818be6a0383c3dcda25dc4e", size = 226202, upload-time = "2026-08-10T22:07:16.942Z" }, +] + [[package]] name = "propcache" version = "0.3.2" @@ -1228,6 +1319,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/32/56/8a7ca5d2cd2cda1d245d34b1c9a942920a718082ae8e54e5f3e5a58b7add/pydantic_core-2.33.2-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:329467cecfb529c925cf2bbd4d60d2c509bc2fb52a20c1045bf09bb70971a9c1", size = 2066757, upload-time = "2025-04-23T18:33:30.645Z" }, ] +[[package]] +name = "pygments" +version = "2.21.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/2e/ced460408999b33da6b31b0021b0f37d329e202d4169aeb164493778f25b/pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c", size = 5005329, upload-time = "2026-08-17T08:02:48.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, +] + +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -1240,6 +1356,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "python-discovery" +version = "1.5.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b2/8f/3c92c45737f654f2488ab3662b7604a55d3d35146d37c9ce80f5c95b95a6/python_discovery-1.5.3.tar.gz", hash = "sha256:e500eb24025fb7c4876c1fdcfbafd9028a10c71b661aee38cb6fb0de594518c1", size = 82477, upload-time = "2026-08-24T14:48:46.396Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/30/12/823d9a321904ccfd2969a24b84fdfd1e6614c707ec569c62879bf1dbc6c5/python_discovery-1.5.3-py3-none-any.whl", hash = "sha256:8305296358f1aa2ed302a25b84be7df84fef8ca47c7dce2da63cb7325333044e", size = 38290, upload-time = "2026-08-24T14:48:45.305Z" }, +] + [[package]] name = "pytz" version = "2025.2" @@ -1499,6 +1627,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a7/c2/fe1e52489ae3122415c51f387e221dd0773709bad6c6cdaa599e8a2c5185/urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc", size = 129795, upload-time = "2025-06-18T14:07:40.39Z" }, ] +[[package]] +name = "virtualenv" +version = "21.7.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "distlib" }, + { name = "filelock" }, + { name = "platformdirs" }, + { name = "python-discovery" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1d/60/fc54e876e34f94dd0cf0185aaecfd4bfa906653f003d9b2fb21428642fca/virtualenv-21.7.5.tar.gz", hash = "sha256:a73c4246fba3c8901ff9717399f466e00eeca5a3834981f1a6ebb4f1e94de2f8", size = 5346743, upload-time = "2026-08-25T05:39:16.14Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/d8/401141bf45637be916c86d325bd821c5838c7eff83294b934cd94e774e4f/virtualenv-21.7.5-py3-none-any.whl", hash = "sha256:e36ca889510ab6cb0b1dca93c59e5431dd4422a3c88f487358d470c90af8c07a", size = 5324697, upload-time = "2026-08-25T05:39:14.229Z" }, +] + [[package]] name = "xxhash" version = "3.5.0"