dataforge-07 0.1.0__py3-none-any.whl

dataforge/bench/methods.py

@@ -0,0 +1,443 @@
+"""Benchmark method implementations for DataForge."""
+
+from __future__ import annotations
+
+import json
+import random
+import time
+from math import ceil
+from statistics import median
+
+from dataforge.bench.core import BenchmarkRepair, SeedBenchmarkResult, quota_units, score_repairs
+from dataforge.bench.groq_client import GroqBenchClient
+from dataforge.datasets.real_world import RealWorldDataset
+from dataforge.detectors import run_all_detectors
+from dataforge.repairers import propose_fixes
+from dataforge.schema_inference import infer_schema
+
+
+def _reproduction_command(method: str, dataset: str, seeds: int) -> str:
+    """Build the canonical reproduction command for one method/dataset selection."""
+    return f"dataforge bench --methods {method} --datasets {dataset} --seeds {seeds}"
+
+
+def _repairs_from_proposed_fixes(dataset: RealWorldDataset) -> list[BenchmarkRepair]:
+    """Run the shipped deterministic detector/repair stack on one dataset."""
+    inferred_schema = infer_schema(dataset.dirty_df.copy(deep=True)).to_schema(
+        include_inferred_constraints=True
+    )
+    issues = run_all_detectors(dataset.dirty_df.copy(deep=True), schema=inferred_schema)
+    proposals = propose_fixes(
+        issues,
+        dataset.dirty_df.copy(deep=True),
+        inferred_schema,
+        cache_dir=None,
+        allow_llm=False,
+    )
+    return [
+        BenchmarkRepair(
+            row=proposal.fix.row,
+            column=proposal.fix.column,
+            new_value=proposal.fix.new_value,
+            reason=proposal.reason,
+        )
+        for proposal in proposals
+    ]
+
+
+def run_heuristic_episode(dataset: RealWorldDataset, *, seed: int) -> SeedBenchmarkResult:
+    """Run the current deterministic DataForge stack as the heuristic baseline."""
+    start = time.perf_counter()
+    repairs = _repairs_from_proposed_fixes(dataset)
+    metrics = score_repairs(dataset.ground_truth, repairs)
+    runtime_s = round(time.perf_counter() - start, 4)
+    return SeedBenchmarkResult(
+        method="heuristic",
+        dataset=dataset.metadata.name,
+        seed=seed,
+        status="ok",
+        precision=metrics.precision,
+        recall=metrics.recall,
+        f1=metrics.f1,
+        tp=metrics.tp,
+        fp=metrics.fp,
+        fn=metrics.fn,
+        avg_steps=float(1 + len(repairs)),
+        llm_calls=0,
+        prompt_tokens=0,
+        completion_tokens=0,
+        quota_units=0.0,
+        runtime_s=runtime_s,
+        provider="local",
+        model="deterministic",
+        reproduction_command=_reproduction_command("heuristic", dataset.metadata.name, 1),
+    )
+
+
+def run_random_episode(dataset: RealWorldDataset, *, seed: int) -> SeedBenchmarkResult:
+    """Run the bounded random baseline on one dataset."""
+    rng = random.Random(seed)
+    start = time.perf_counter()
+    budget = min(200, max(25, ceil(len(dataset.ground_truth) / 10)))
+    column_values = {
+        column: [str(value) for value in dataset.dirty_df[column].tolist()]
+        for column in dataset.canonical_columns
+    }
+    repairs: list[BenchmarkRepair] = []
+    for _ in range(budget):
+        row_index = rng.randrange(len(dataset.dirty_df.index))
+        column = rng.choice(dataset.canonical_columns)
+        new_value = rng.choice(column_values[column])
+        repairs.append(
+            BenchmarkRepair(
+                row=row_index,
+                column=column,
+                new_value=new_value,
+                reason="random baseline",
+            )
+        )
+    metrics = score_repairs(dataset.ground_truth, repairs)
+    runtime_s = round(time.perf_counter() - start, 4)
+    return SeedBenchmarkResult(
+        method="random",
+        dataset=dataset.metadata.name,
+        seed=seed,
+        status="ok",
+        precision=metrics.precision,
+        recall=metrics.recall,
+        f1=metrics.f1,
+        tp=metrics.tp,
+        fp=metrics.fp,
+        fn=metrics.fn,
+        avg_steps=float(budget),
+        llm_calls=0,
+        prompt_tokens=0,
+        completion_tokens=0,
+        quota_units=0.0,
+        runtime_s=runtime_s,
+        provider="local",
+        model="random",
+        reproduction_command=_reproduction_command("random", dataset.metadata.name, 1),
+    )
+
+
+def _chunk_records(dataset: RealWorldDataset, row_indices: tuple[int, ...]) -> list[dict[str, str]]:
+    """Serialize one row chunk for prompting."""
+    records: list[dict[str, str]] = []
+    for row_index in row_indices:
+        row_payload: dict[str, str] = {"_row": str(row_index)}
+        for column in dataset.canonical_columns:
+            row_payload[column] = str(dataset.dirty_df.iloc[row_index][column])
+        records.append(row_payload)
+    return records
+
+
+def _column_stats(
+    dataset: RealWorldDataset, columns: list[str]
+) -> dict[str, dict[str, str | float | int]]:
+    """Return simple benchmark-local column statistics for ReAct prompting."""
+    stats: dict[str, dict[str, str | float | int]] = {}
+    for column in columns:
+        series = dataset.dirty_df[column].astype(str)
+        non_empty = [value for value in series.tolist() if value != ""]
+        numeric_values: list[float] = []
+        for value in non_empty:
+            try:
+                numeric_values.append(float(value))
+            except ValueError:
+                continue
+        stats[column] = {
+            "non_empty_count": len(non_empty),
+            "unique_count": len(set(non_empty)),
+        }
+        if numeric_values:
+            stats[column]["median"] = round(float(median(numeric_values)), 4)
+    return stats
+
+
+def _strip_json_line_comments(text: str) -> str:
+    """Remove JavaScript-style line comments outside JSON strings."""
+    result: list[str] = []
+    in_string = False
+    escaped = False
+    index = 0
+    while index < len(text):
+        char = text[index]
+        next_char = text[index + 1] if index + 1 < len(text) else ""
+        if in_string:
+            result.append(char)
+            if escaped:
+                escaped = False
+            elif char == "\\":
+                escaped = True
+            elif char == '"':
+                in_string = False
+            index += 1
+            continue
+        if char == '"':
+            in_string = True
+            result.append(char)
+            index += 1
+            continue
+        if char == "/" and next_char == "/":
+            index += 2
+            while index < len(text) and text[index] not in "\r\n":
+                index += 1
+            continue
+        result.append(char)
+        index += 1
+    return "".join(result)
+
+
+def _extract_json_object(text: str) -> dict[str, object] | None:
+    """Parse the first JSON object found in an LLM response string."""
+    stripped = text.strip()
+    if stripped.startswith("```"):
+        stripped = stripped.strip("`")
+        if stripped.lower().startswith("json"):
+            stripped = stripped[4:].strip()
+    stripped = _strip_json_line_comments(stripped)
+    decoder = json.JSONDecoder()
+    for offset, char in enumerate(stripped):
+        if char != "{":
+            continue
+        try:
+            payload, _ = decoder.raw_decode(stripped[offset:])
+        except json.JSONDecodeError:
+            continue
+        if isinstance(payload, dict):
+            return payload
+    return None
+
+
+def _repairs_from_payload(payload: dict[str, object]) -> list[BenchmarkRepair]:
+    """Convert a parsed JSON payload into benchmark repairs."""
+    raw_repairs = payload.get("repairs", [])
+    if not isinstance(raw_repairs, list):
+        return []
+    repairs: list[BenchmarkRepair] = []
+    for raw_repair in raw_repairs:
+        if not isinstance(raw_repair, dict):
+            continue
+        row = raw_repair.get("row")
+        column = raw_repair.get("column")
+        new_value = raw_repair.get("new_value")
+        reason = raw_repair.get("reason", "LLM repair")
+        if (
+            not isinstance(row, int)
+            or not isinstance(column, str)
+            or not isinstance(new_value, str)
+        ):
+            continue
+        repairs.append(
+            BenchmarkRepair(
+                row=row,
+                column=column,
+                new_value=new_value,
+                reason=str(reason),
+            )
+        )
+    return repairs
+
+
+def run_llm_zeroshot_episode(
+    dataset: RealWorldDataset,
+    *,
+    seed: int,
+    client: GroqBenchClient,
+) -> SeedBenchmarkResult:
+    """Run the zero-shot Groq baseline across fixed contiguous row chunks."""
+    start = time.perf_counter()
+    llm_calls = 0
+    prompt_tokens = 0
+    completion_tokens = 0
+    warnings: list[str] = []
+    repairs: list[BenchmarkRepair] = []
+
+    for row_indices in chunk_row_indices(len(dataset.dirty_df.index)):
+        chunk_payload = _chunk_records(dataset, row_indices)
+        messages = [
+            {
+                "role": "system",
+                "content": (
+                    "You are benchmarking tabular data cleaning. Reply with strict JSON: "
+                    '{"repairs":[{"row":0,"column":"Column","new_value":"value","reason":"why"}]}.'
+                ),
+            },
+            {
+                "role": "user",
+                "content": json.dumps(
+                    {
+                        "dataset": dataset.metadata.name,
+                        "columns": list(dataset.canonical_columns),
+                        "rows": chunk_payload,
+                    },
+                    sort_keys=True,
+                ),
+            },
+        ]
+        completion = client.complete(messages)
+        llm_calls += 1
+        prompt_tokens += completion.prompt_tokens
+        completion_tokens += completion.completion_tokens
+        warnings.extend(list(completion.warnings))
+        parsed = _extract_json_object(completion.text)
+        if parsed is not None:
+            repairs.extend(_repairs_from_payload(parsed))
+
+    metrics = score_repairs(dataset.ground_truth, repairs)
+    runtime_s = round(time.perf_counter() - start, 4)
+    return SeedBenchmarkResult(
+        method="llm_zeroshot",
+        dataset=dataset.metadata.name,
+        seed=seed,
+        status="ok",
+        precision=metrics.precision,
+        recall=metrics.recall,
+        f1=metrics.f1,
+        tp=metrics.tp,
+        fp=metrics.fp,
+        fn=metrics.fn,
+        avg_steps=float(llm_calls),
+        llm_calls=llm_calls,
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        quota_units=quota_units(
+            llm_calls=llm_calls,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+        ),
+        runtime_s=runtime_s,
+        provider="groq",
+        model=client.model,
+        warnings=warnings,
+        reproduction_command=_reproduction_command("llm_zeroshot", dataset.metadata.name, 1),
+    )
+
+
+def run_llm_react_episode(
+    dataset: RealWorldDataset,
+    *,
+    seed: int,
+    client: GroqBenchClient,
+) -> SeedBenchmarkResult:
+    """Run the constrained ReAct-style Groq baseline with one optional tool step."""
+    start = time.perf_counter()
+    llm_calls = 0
+    tool_calls = 0
+    prompt_tokens = 0
+    completion_tokens = 0
+    warnings: list[str] = []
+    repairs: list[BenchmarkRepair] = []
+
+    for row_indices in chunk_row_indices(len(dataset.dirty_df.index)):
+        chunk_payload = _chunk_records(dataset, row_indices)
+        schema_summary = {
+            "dataset": dataset.metadata.name,
+            "columns": list(dataset.canonical_columns),
+            "chunk_rows": len(row_indices),
+        }
+        messages = [
+            {
+                "role": "system",
+                "content": (
+                    "You are benchmarking tabular data cleaning with a constrained tool loop. "
+                    "Respond with one JSON action object. Allowed actions: "
+                    "inspect_rows, column_stats, submit_repairs, finish."
+                ),
+            },
+            {
+                "role": "user",
+                "content": json.dumps(
+                    {
+                        "schema_summary": schema_summary,
+                        "rows": chunk_payload,
+                    },
+                    sort_keys=True,
+                ),
+            },
+        ]
+        first = client.complete(messages)
+        llm_calls += 1
+        prompt_tokens += first.prompt_tokens
+        completion_tokens += first.completion_tokens
+        warnings.extend(list(first.warnings))
+        first_payload = _extract_json_object(first.text)
+        if first_payload is None:
+            continue
+
+        action = first_payload.get("action")
+        if action == "submit_repairs":
+            repairs.extend(_repairs_from_payload(first_payload))
+            continue
+        if action == "finish":
+            continue
+
+        tool_result: dict[str, object]
+        if action == "inspect_rows":
+            requested_rows = first_payload.get("row_indices", [])
+            if not isinstance(requested_rows, list):
+                requested_rows = []
+            safe_rows = [
+                row for row in requested_rows if isinstance(row, int) and row in row_indices
+            ]
+            tool_result = {"rows": _chunk_records(dataset, tuple(safe_rows))}
+        elif action == "column_stats":
+            requested_columns = first_payload.get("columns", [])
+            if not isinstance(requested_columns, list):
+                requested_columns = []
+            safe_columns = [
+                column
+                for column in requested_columns
+                if isinstance(column, str) and column in dataset.canonical_columns
+            ]
+            tool_result = {"column_stats": _column_stats(dataset, safe_columns)}
+        else:
+            continue
+        tool_calls += 1
+        messages.append({"role": "assistant", "content": first.text})
+        messages.append({"role": "user", "content": json.dumps(tool_result, sort_keys=True)})
+        second = client.complete(messages)
+        llm_calls += 1
+        prompt_tokens += second.prompt_tokens
+        completion_tokens += second.completion_tokens
+        warnings.extend(list(second.warnings))
+        second_payload = _extract_json_object(second.text)
+        if second_payload is not None and second_payload.get("action") == "submit_repairs":
+            repairs.extend(_repairs_from_payload(second_payload))
+
+    metrics = score_repairs(dataset.ground_truth, repairs)
+    runtime_s = round(time.perf_counter() - start, 4)
+    return SeedBenchmarkResult(
+        method="llm_react",
+        dataset=dataset.metadata.name,
+        seed=seed,
+        status="ok",
+        precision=metrics.precision,
+        recall=metrics.recall,
+        f1=metrics.f1,
+        tp=metrics.tp,
+        fp=metrics.fp,
+        fn=metrics.fn,
+        avg_steps=float(llm_calls + tool_calls),
+        llm_calls=llm_calls,
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        quota_units=quota_units(
+            llm_calls=llm_calls,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+        ),
+        runtime_s=runtime_s,
+        provider="groq",
+        model=client.model,
+        warnings=warnings,
+        reproduction_command=_reproduction_command("llm_react", dataset.metadata.name, 1),
+    )
+
+
+def chunk_row_indices(n_rows: int) -> tuple[tuple[int, ...], ...]:
+    """Local import wrapper that avoids circular imports in the LLM helpers."""
+    from dataforge.bench.core import chunk_row_indices as _chunk_row_indices
+
+    return _chunk_row_indices(n_rows)