Speed up dedup/embedding/scoring, add local-model support and CLI info/sample commands

Author: voidstackloop

benchmarks/benchmark_buffdata.py

@@ -975,6 +975,7 @@ async def main(args: argparse.Namespace) -> None:
     parser.add_argument("--datasets", nargs="+", choices=sorted(DATASETS))
     parser.add_argument("--gemini-audit-rows", type=int, default=0)
     parser.add_argument("--gemini-model", default="gemini-3.7-flash")
+    parser.add_argument("--openai-model", default="gpt-5")
     parser.add_argument("--anthropic-model", default="claude-sonnet-5")
-    parser.add_argument("--audit-provider", choices=["gemini", "anthropic", "both"], default="gemini")
+    parser.add_argument("--audit-provider", choices=["gemini", "openai", "anthropic", "both"], default="gemini")
     asyncio.run(main(parser.parse_args()))

benchmarks/benchmark_claude.py

@@ -46,6 +46,12 @@
     --provider openai_compatible --base-url https://proxy.corp/v1 \\
     --model claude-sonnet-5
 
+  # local model via LM Studio on a shared box (free inference, so audit +
+  # classification can stay on); --model must match the id LM Studio shows
+  python benchmarks/benchmark_claude.py --smoke \\
+    --provider lmstudio --base-url http://bionic:1234/v1 \\
+    --model qwen3-27b
+
   # offline plumbing check: zero API calls
   python benchmarks/benchmark_claude.py --smoke --mock
 
@@ -75,6 +81,18 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 
+# Parse --request-timeout BEFORE importing the client stack: engine/client.py
+# binds DEFAULT_REQUEST_TIMEOUT_SECONDS at import time, so the env var must be
+# set first. Slow local models (27B reasoning on CPU/offload) need 900-1800s;
+# cloud APIs are fine with the 120s default.
+_request_parser = argparse.ArgumentParser(add_help=False)
+_request_parser.add_argument("--request-timeout", type=float, default=None)
+_request_known, _ = _request_parser.parse_known_args()
+if _request_known.request_timeout:
+    import os as _os
+
+    _os.environ["BUFFDATA_REQUEST_TIMEOUT"] = str(_request_known.request_timeout)
+
 from benchmark_buffdata import (  # noqa: E402
     build_vocab,
     make_dirty,
@@ -90,24 +108,41 @@
 DEFAULT_DATASETS = ["ag_news", "imdb"]
 
 
+# provider -> (base-url env var, api-key env var) read by create_llm_client
+# when --base-url is not passed explicitly. Local servers (ollama/lmstudio/...)
+# also have loopback presets, so --base-url is only needed for shared boxes
+# like --base-url http://bionic:1234/v1.
+_GATEWAY_ENVS = {
+    "openai_compatible": ("OPENAI_COMPATIBLE_BASE_URL", "OPENAI_COMPATIBLE_API_KEY"),
+    "ollama": ("OLLAMA_BASE_URL", "OLLAMA_API_KEY"),
+    "lmstudio": ("LMSTUDIO_BASE_URL", "LMSTUDIO_API_KEY"),
+    "vllm": ("VLLM_BASE_URL", "VLLM_API_KEY"),
+    "llamacpp": ("LLAMACPP_BASE_URL", "LLAMACPP_API_KEY"),
+}
+
+
 def _configure_gateway(provider: str, base_url: str | None) -> None:
     """Forward gateway settings into the env vars the client layer reads.
 
     The optimize() helper below builds its own PipelineConfig without a
-    base_url, but create_llm_client(provider="openai_compatible") falls back
-    to OPENAI_COMPATIBLE_BASE_URL / OPENAI_COMPATIBLE_API_KEY env vars -- so
-    exporting here makes both the direct classifier calls and the full
-    pipeline audit path reach the proxy. Reuses .env ANTHROPIC_API_KEY as the
-    gateway key when no gateway-specific key is set.
+    base_url, but create_llm_client falls back to the provider's *_BASE_URL /
+    *_API_KEY env vars -- so exporting here makes both the direct classifier
+    calls and the full pipeline audit path reach the proxy or local server.
+    For openai_compatible, reuses .env ANTHROPIC_API_KEY as the gateway key
+    when no gateway-specific key is set. Local servers need no key.
     """
     import os
 
+    envs = _GATEWAY_ENVS.get(provider)
+    if not envs:
+        return
+    base_env, key_env = envs
     if base_url:
-        os.environ["OPENAI_COMPATIBLE_BASE_URL"] = base_url
-    if provider == "openai_compatible" and not os.getenv("OPENAI_COMPATIBLE_API_KEY"):
+        os.environ[base_env] = base_url
+    if provider == "openai_compatible" and not os.getenv(key_env):
         fallback = os.getenv("ANTHROPIC_API_KEY")
         if fallback:
-            os.environ["OPENAI_COMPATIBLE_API_KEY"] = fallback
+            os.environ[key_env] = fallback
 
 
 def _synthetic_rows(count: int, classes: int, seed: int) -> list[dict[str, Any]]:
@@ -225,6 +260,25 @@ async def measure_classification(
     }
 
 
+async def _safe_optimize(
+    rows: list[dict[str, Any]], label: str, **kwargs: Any,
+) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+    """optimize() that never kills the benchmark: on provider failure (wedged
+    local server, dead quota, 503 spike) fall back to the unoptimized input so
+    recovery still reports -- with the error recorded in quality instead of
+    silently presenting fallback numbers as a real result."""
+    try:
+        return await optimize(rows, **kwargs)
+    except Exception as exc:  # noqa: BLE001
+        print(f"  [{label}] optimize failed, using unoptimized input: {exc}", flush=True)
+        return list(rows), {
+            "error": f"{type(exc).__name__}: {exc}",
+            "input_rows": len(rows), "output_rows": len(rows),
+            "provider": kwargs.get("provider"), "model": kwargs.get("anthropic_model"),
+            "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
+        }
+
+
 async def run_smoke(args: argparse.Namespace) -> dict[str, Any]:
     """Synthetic end-to-end check without Hugging Face downloads."""
     print("=== smoke: synthetic binary classification (no HF download) ===", flush=True)
@@ -239,9 +293,11 @@ async def run_smoke(args: argparse.Namespace) -> dict[str, Any]:
     # force the audit off (quality_mode="off") instead of attempting a live
     # sampled audit with an absent key.
     audit_rows = 0 if args.mock else args.audit_rows
-    clean_opt, clean_q = await optimize(clean, provider=args.provider, anthropic_model=args.model)
-    dirty_opt, dirty_q = await optimize(
-        optimizer_input,
+    clean_opt, clean_q = await _safe_optimize(
+        clean, "smoke-clean", provider=args.provider, anthropic_model=args.model,
+    )
+    dirty_opt, dirty_q = await _safe_optimize(
+        optimizer_input, "smoke-dirty",
         gemini_audit_rows=audit_rows,
         gemini_model=args.model,
         anthropic_model=args.model,
@@ -304,9 +360,11 @@ async def run_dataset(
     clean_test = stratified_rows(test_split, args.test_rows, classes, text_fields, 200 + dataset_offset, label_field=label_col)
     _, optimizer_input, defects = make_dirty(clean_train, classes, 300 + dataset_offset)
 
-    clean_opt, clean_q = await optimize(clean_train, provider=args.provider, anthropic_model=args.model)
-    dirty_opt, dirty_q = await optimize(
-        optimizer_input,
+    clean_opt, clean_q = await _safe_optimize(
+        clean_train, f"{name}-clean", provider=args.provider, anthropic_model=args.model,
+    )
+    dirty_opt, dirty_q = await _safe_optimize(
+        optimizer_input, f"{name}-dirty",
         gemini_audit_rows=args.audit_rows,
         gemini_model=args.model,
         anthropic_model=args.model,
@@ -367,7 +425,7 @@ def markdown_report(payload: dict[str, Any]) -> str:
         f"# {title}",
         "",
         f"Provider: `{m['provider']}` | Model: `{m['model']}` | Base URL: `{m.get('base_url') or '-'}` | Seeds: {m['seeds']} | Epochs: {m['epochs']} |",
-        f"Audit rows: {m['audit_rows']} | Classify sample: {m['classify_sample_size']}",
+        f"Audit rows: {m['audit_rows']} | Classify sample: {m['classify_sample_size']} | Request timeout: {m.get('request_timeout_s', '-')}s",
         "",
     ]
     if payload.get("smoke"):
@@ -396,9 +454,12 @@ def markdown_report(payload: dict[str, Any]) -> str:
         q = r.get("dirty_audit_quality", {})
         u = q.get("usage", {})
         cl = (r.get("classification") or {}).get("usage", {})
+        audit_cell: Any = q.get("audit_avg_score", "n/a")
+        if q.get("error"):
+            audit_cell = f"error ({str(q['error'])[:80]})"
         lines.append(
             f"| {r['dataset']} | {u.get('input_tokens', 0)}/{u.get('output_tokens', 0)}/{u.get('total_tokens', 0)} | "
-            f"{q.get('audit_avg_score', 'n/a')} | {cl.get('total_tokens', cl.get('total', 0))} |"
+            f"{audit_cell} | {cl.get('total_tokens', cl.get('total', 0))} |"
         )
     lines.append("")
     return "\n".join(lines)
@@ -423,9 +484,13 @@ async def main(args: argparse.Namespace) -> None:
 
     from benchmark_buffdata import DATASETS
 
+    from buffdata.engine.client import DEFAULT_REQUEST_TIMEOUT_SECONDS
+
+    print(f"request timeout: {DEFAULT_REQUEST_TIMEOUT_SECONDS:g}s", flush=True)
     payload: dict[str, Any] = {
         "method": {
             "provider": args.provider, "model": args.model, "base_url": args.base_url,
+            "request_timeout_s": DEFAULT_REQUEST_TIMEOUT_SECONDS,
             "seeds": args.seeds,
             "epochs": args.epochs, "audit_rows": args.audit_rows,
             "classify_sample_size": args.classify_sample_size,
@@ -465,4 +530,9 @@ async def main(args: argparse.Namespace) -> None:
     parser.add_argument("--skip-classification", action="store_true")
     parser.add_argument("--smoke", action="store_true", help="Synthetic data, no HF download.")
     parser.add_argument("--mock", action="store_true", help="Zero API calls (offline plumbing check).")
+    parser.add_argument(
+        "--request-timeout", type=float, default=None,
+        help="Per-request seconds (sets BUFFDATA_REQUEST_TIMEOUT before client import). "
+        "Use 900-1800 for slow local models; default 120 suits cloud APIs.",
+    )
     asyncio.run(main(parser.parse_args()))

benchmarks/benchmark_claude_classification.py

@@ -50,6 +50,16 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 
+# Parse --request-timeout BEFORE any client import (engine/client.py binds its
+# default timeout at import time). Slow local models need 900-1800s.
+_request_parser = argparse.ArgumentParser(add_help=False)
+_request_parser.add_argument("--request-timeout", type=float, default=None)
+_request_known, _ = _request_parser.parse_known_args()
+if _request_known.request_timeout:
+    import os as _os
+
+    _os.environ["BUFFDATA_REQUEST_TIMEOUT"] = str(_request_known.request_timeout)
+
 DEFAULT_PROVIDER = "anthropic"
 DEFAULT_MODEL = "claude-sonnet-5"
 DEFAULT_DATASETS = ["imdb", "ag_news", "emotion"]
@@ -106,23 +116,39 @@ async def generate_structured_async(self, prompt: str, response_schema, model=No
         )
 
 
+# provider -> (base-url env var, api-key env var) read by create_llm_client
+# when --base-url is not passed explicitly (e.g. --base-url http://bionic:1234/v1).
+_GATEWAY_ENVS = {
+    "openai_compatible": ("OPENAI_COMPATIBLE_BASE_URL", "OPENAI_COMPATIBLE_API_KEY"),
+    "ollama": ("OLLAMA_BASE_URL", "OLLAMA_API_KEY"),
+    "lmstudio": ("LMSTUDIO_BASE_URL", "LMSTUDIO_API_KEY"),
+    "vllm": ("VLLM_BASE_URL", "VLLM_API_KEY"),
+    "llamacpp": ("LLAMACPP_BASE_URL", "LLAMACPP_API_KEY"),
+}
+
+
 def _configure_gateway(provider: str, base_url: str | None) -> None:
     """Forward gateway settings into the env vars the client layer reads.
 
     create_llm_client(provider="openai_compatible") resolves its endpoint from
     --base-url or OPENAI_COMPATIBLE_BASE_URL and its key from
-    OPENAI_COMPATIBLE_API_KEY. Reuse the existing .env ANTHROPIC_API_KEY when
-    the gateway-specific key is unset so external-proxy Claude keys work
-    without copying secrets by hand.
+    OPENAI_COMPATIBLE_API_KEY (same pattern for the local-server providers).
+    Reuse the existing .env ANTHROPIC_API_KEY when the gateway-specific key
+    is unset so external-proxy Claude keys work without copying secrets by
+    hand. Local servers need no key.
     """
     import os
 
+    envs = _GATEWAY_ENVS.get(provider)
+    if not envs:
+        return
+    base_env, key_env = envs
     if base_url:
-        os.environ["OPENAI_COMPATIBLE_BASE_URL"] = base_url
-    if provider == "openai_compatible" and not os.getenv("OPENAI_COMPATIBLE_API_KEY"):
+        os.environ[base_env] = base_url
+    if provider == "openai_compatible" and not os.getenv(key_env):
         fallback = os.getenv("ANTHROPIC_API_KEY")
         if fallback:
-            os.environ["OPENAI_COMPATIBLE_API_KEY"] = fallback
+            os.environ[key_env] = fallback
 
 
 async def classify_and_score(
@@ -235,6 +261,9 @@ async def main(args: argparse.Namespace) -> None:
     results_path = output_dir / "results.json"
     report_path = output_dir / "REPORT.md"
 
+    from buffdata.engine.client import DEFAULT_REQUEST_TIMEOUT_SECONDS
+
+    print(f"request timeout: {DEFAULT_REQUEST_TIMEOUT_SECONDS:g}s", flush=True)
     results: list[dict[str, Any]] = []
     if args.smoke:
         print("=== smoke: synthetic binary + multi-class (no downloads) ===", flush=True)
@@ -285,4 +314,9 @@ async def main(args: argparse.Namespace) -> None:
     parser.add_argument("--sample-size", type=int, default=20)
     parser.add_argument("--smoke", action="store_true")
     parser.add_argument("--mock", action="store_true")
+    parser.add_argument(
+        "--request-timeout", type=float, default=None,
+        help="Per-request seconds (sets BUFFDATA_REQUEST_TIMEOUT before client import). "
+        "Use 900-1800 for slow local models; default 120 suits cloud APIs.",
+    )
     asyncio.run(main(parser.parse_args()))

buffdata/cli/info.py

@@ -0,0 +1,113 @@
+"""Quick dataset inspection commands for buffdata CLI."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional
+
+import typer
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+from buffdata.models.formats import _is_cloud_url, iter_dataset, read_dataset
+
+console = Console()
+
+
+def info_cmd(
+    input_file: str = typer.Argument(..., help="Path to dataset, or a cloud URL"),
+    max_rows: Optional[int] = typer.Option(None, "--max-rows", "-n", help="Only inspect first N rows"),
+):
+    """Show quick dataset info: format, row count, columns, and sample values."""
+    from buffdata.models.formats import _path_format, _cloud_format
+
+    if _is_cloud_url(input_file):
+        dataset_format = _cloud_format(input_file)
+        source = input_file
+    else:
+        path = Path(input_file)
+        if not path.exists():
+            console.print(f"[red]File not found: {path}[/red]")
+            raise typer.Exit(1)
+        dataset_format = _path_format(path)
+        source = str(path)
+
+    items = []
+    row_count = 0
+    columns = set()
+    sample_values = {}
+
+    for item in iter_dataset(input_file):
+        row_count += 1
+        if row_count <= 5:
+            items.append(item)
+            d = item.to_dict()
+            columns.update(d.keys())
+            for k, v in d.items():
+                if k not in sample_values:
+                    sample_values[k] = v
+        if max_rows and row_count >= max_rows:
+            break
+
+    table = Table(title="Dataset Info", border_style="cyan")
+    table.add_column("Property", style="bold")
+    table.add_column("Value")
+    table.add_row("Source", source)
+    table.add_row("Format", dataset_format)
+    table.add_row("Rows", str(row_count) + (" (limited)" if max_rows else ""))
+    table.add_row("Columns", ", ".join(sorted(columns)) if columns else "-")
+    console.print(table)
+
+    if items:
+        console.print()
+        preview_table = Table(title="Preview (first rows)", border_style="green")
+        preview_cols = sorted(columns)[:6]
+        for col in preview_cols:
+            preview_table.add_column(col, max_width=40)
+        for item in items[:3]:
+            d = item.to_dict()
+            row = [str(d.get(col, ""))[:40] for col in preview_cols]
+            preview_table.add_row(*row)
+        console.print(preview_table)
+
+
+def sample_cmd(
+    input_file: str = typer.Argument(..., help="Path to dataset, or a cloud URL"),
+    n: int = typer.Option(5, "-n", "--count", help="Number of rows to show"),
+    offset: int = typer.Option(0, "--offset", "-o", help="Skip first N rows"),
+    format: str = typer.Option("table", "--format", "-f", help="Output format: table, json, jsonl"),
+):
+    """Preview N rows from a dataset."""
+    import json
+
+    items = []
+    for i, item in enumerate(iter_dataset(input_file)):
+        if i < offset:
+            continue
+        if len(items) >= n:
+            break
+        items.append(item)
+
+    if not items:
+        console.print("[yellow]No rows found.[/yellow]")
+        return
+
+    if format == "json":
+        console.print_json(data=[item.to_dict() for item in items])
+    elif format == "jsonl":
+        for item in items:
+            console.print(json.dumps(item.to_dict(), ensure_ascii=False))
+    else:
+        table = Table(title=f"Sample ({len(items)} rows)", border_style="green")
+        all_keys = set()
+        for item in items:
+            all_keys.update(item.to_dict().keys())
+        cols = sorted(all_keys)[:5]
+        for col in cols:
+            table.add_column(col, max_width=50)
+        for item in items:
+            d = item.to_dict()
+            table.add_row(*[str(d.get(col, ""))[:50] for col in cols])
+        console.print(table)