weco 0.3.0__tar.gz → 0.3.2__tar.gz

{weco-0.3.0 → weco-0.3.2}/.gitignore

@@ -84,3 +84,8 @@ repomix-output.*
 
 # Claude config
 .claude/
+
+# Example: extract-line-plot generated artifacts
+examples/extract-line-plot/predictions/
+examples/extract-line-plot/subset_line_*/
+examples/extract-line-plot/subset_line_*.zip

{weco-0.3.0 → weco-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.3.0
+Version: 0.3.2
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: 
@@ -219,6 +219,7 @@ Requires-Dist: packaging
 Requires-Dist: gitingest
 Requires-Dist: fastapi
 Requires-Dist: slowapi
+Requires-Dist: psutil
 Provides-Extra: dev
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: build; extra == "dev"

{weco-0.3.0 → weco-0.3.2}/examples/cuda/README.md

@@ -7,9 +7,11 @@ This approach aims for low-level optimization beyond standard PyTorch or even Tr
 
 Install the CLI and dependencies for the example:
 ```bash
-pip install weco torch ninja triton
+pip install weco ninja numpy torch triton
 ```
-> **Note:** This example requires a compatible NVIDIA GPU and the CUDA Toolkit installed on your system for compiling and running the generated CUDA code.
+> **Note:**
+> 1. This example requires a compatible NVIDIA GPU and the CUDA Toolkit installed on your system for compiling and running the generated CUDA code.
+> 2. If compatible, install [flash attention](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#installation-and-features) (`pip install flash-attn --no-build-isolation`).
 
 ## Run Weco
 
@@ -20,8 +22,9 @@ weco run --source optimize.py \
      --metric speedup \
      --goal maximize \
      --steps 50 \
-     --model o4-mini \
-     --additional-instructions "Write in-line CUDA using pytorch's load_inline() to optimize the code while ensuring a small max float diff. Maintain the same code format. Do not use any fallbacks. Assume any required dependencies are installed and data is already on the gpu."
+     --model gpt-5 \
+     --additional-instructions "Write in-line CUDA using pytorch's load_inline() to optimize the code while ensuring a small max float diff. Maintain the same code interface. Do not use any fallbacks and never use the build_directory arg for load_inline(). Assume any required dependencies are installed and data is already on the gpu." \
+     --eval-timeout 600
 ```
 
 ### Explanation
@@ -31,8 +34,9 @@ weco run --source optimize.py \
 *   `--metric speedup`: The optimization target metric.
 *   `--goal maximize`: Weco aims to increase the speedup.
 *   `--steps 50`: The number of optimization iterations.
-*   `--model o4-mini`: The LLM used for code generation.
+*   `--model gpt-5`: The LLM used for code generation.
 *   `--additional-instructions "..."`: Provides guidance to the LLM on the optimization approach.
+*   `--eval-timeout 600`: Stop runnning the evaluation script if it does not complete in 600 seconds.
 
 Weco will iteratively modify `optimize.py`, generating and integrating CUDA C++ code, guided by the evaluation results and the additional instructions provided.
 

{weco-0.3.0 → weco-0.3.2}/examples/cuda/evaluate.py

@@ -1,5 +1,6 @@
 import sys
 import os
+import shutil
 import pathlib
 import importlib
 import importlib.util
@@ -55,10 +56,9 @@ class Model(nn.Module):
 
 
 ########################################################
-# Weco Solution
+# Benchmark
 ########################################################
 def load_module_from_path(module_path: str, add_to_sys_modules: bool = False):
-    # Clean out all old compiled extensions to prevent namespace collisions during build
     module_path = pathlib.Path(module_path)
     name = module_path.stem
     spec = importlib.util.spec_from_file_location(name, module_path)
@@ -69,12 +69,6 @@ def load_module_from_path(module_path: str, add_to_sys_modules: bool = False):
     return mod
 
 
-########################################################
-# Benchmark
-########################################################
-os.environ["MAX_JOBS"] = "1"  # number of workers for building with ninja
-
-
 def get_inputs(batch_size, seq_len, n_embd, device):
     return torch.randn(batch_size, seq_len, n_embd, device=device, dtype=torch.float32)
 
@@ -86,6 +80,12 @@ if __name__ == "__main__":
     parser.add_argument("--solution-path", type=str, required=True)
     args = parser.parse_args()
 
+    # setup local cache for PyTorch extensions
+    cache_dir = pathlib.Path.cwd() / ".weco-temp/torch_extensions"
+    shutil.rmtree(cache_dir.parent, ignore_errors=True)
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    os.environ["TORCH_EXTENSIONS_DIR"] = str(cache_dir)
+
     # benchmarking parameters
     n_correctness_trials = 10
     correctness_tolerance = 1e-5
@@ -145,3 +145,6 @@ if __name__ == "__main__":
     t_avg_optimized = do_bench(lambda: solution_model(inputs), warmup=warmup_ms, rep=rep_ms)
     print(f"optimized time: {t_avg_optimized:.2f}ms")
     print(f"speedup: {t_avg_baseline / t_avg_optimized:.2f}x")
+
+    # clean up
+    shutil.rmtree(cache_dir.parent, ignore_errors=True)

weco-0.3.2/examples/extract-line-plot/README.md

@@ -0,0 +1,72 @@
+## Extract Line Plot (Chart → CSV) with a VLM
+
+This example is about optimizing an AI feature that turns image of chart into a table in csv format.
+
+### Prerequisites
+
+- Python 3.9+
+- `uv` installed (see `https://docs.astral.sh/uv/`)
+- An OpenAI API key in your environment:
+
+```bash
+export OPENAI_API_KEY=your_key_here
+```
+
+### Files
+
+- `prepare_data.py`: downloads ChartQA (full) and prepares a 100-sample subset of line charts.
+- `optimize.py`: baseline VLM function (`VLMExtractor.image_to_csv`) to be optimized.
+- `eval.py`: evaluation harness that runs the baseline on images and reports a similarity score as "accuracy".
+
+Generated artifacts (gitignored):
+- `subset_line_100/` and `subset_line_100.zip`
+- `predictions/`
+
+### 1) Prepare the data
+
+From the repo root or this directory:
+
+```bash
+cd examples/extract-line-plot
+uv run --with huggingface_hub python prepare_data.py
+```
+
+Notes:
+- Downloads the ChartQA dataset snapshot and auto-extracts `ChartQA Dataset.zip` if needed.
+- Produces `subset_line_100/` with `index.csv`, `images/`, and `tables/`.
+
+### 2) Run a baseline evaluation once
+
+```bash
+uv run --with openai python eval.py --max-samples 10 --num-workers 4
+```
+
+This writes predicted CSVs to `predictions/` and prints a final line like `accuracy: 0.32`.
+
+Metric definition (summarized):
+- Per-sample score = 0.2 × header match + 0.8 × Jaccard(similarity of content rows).
+- Reported `accuracy` is the mean score over all evaluated samples.
+
+### 3) Optimize the baseline with Weco
+
+Run Weco to iteratively improve `optimize.py` using 100 examples and many workers:
+
+```bash
+weco run --source optimize.py --eval-command 'uv run --with openai python eval.py --max-samples 100 --num-workers 50' --metric accuracy --goal maximize --steps 20 --model gpt-5
+```
+
+Arguments:
+- `--source optimize.py`: file that Weco will edit to improve results.
+- `--eval-command '…'`: command Weco executes to measure the metric.
+- `--metric accuracy`: Weco parses `accuracy: <value>` from `eval.py` output.
+- `--goal maximize`: higher is better.
+- `--steps 20`: number of optimization iterations.
+- `--model gpt-5`: model used by Weco to propose edits (change as desired).
+
+### Tips
+
+- Ensure your OpenAI key has access to a vision-capable model (default: `gpt-4o-mini` in the eval; change via `--model`).
+- Adjust `--num-workers` to balance throughput and rate limits.
+- You can tweak baseline behavior in `optimize.py` (prompt, temperature) — Weco will explore modifications automatically during optimization.
+
+

weco-0.3.2/examples/extract-line-plot/eval.py

@@ -0,0 +1,370 @@
+import argparse
+import csv
+import math
+import os
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+from optimize import VLMExtractor
+
+try:
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+except Exception:  # pragma: no cover - optional dependency
+    plt = None
+
+
+def read_index(index_csv_path: Path) -> List[Tuple[str, Path, Path]]:
+    rows: List[Tuple[str, Path, Path]] = []
+    with open(index_csv_path, "r", encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            rows.append((row["id"].strip(), Path(row["image"].strip()), Path(row["table"].strip())))
+    return rows
+
+
+def write_csv(output_dir: Path, example_id: str, csv_text: str) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    out_path = output_dir / f"{example_id}.csv"
+    out_path.write_text(csv_text, encoding="utf-8")
+    return out_path
+
+
+def read_csv_table(path: Path) -> Tuple[List[str], List[List[str]]]:
+    header: List[str] = []
+    rows: List[List[str]] = []
+    with open(path, "r", encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for row in reader:
+            cleaned = [cell.strip() for cell in row]
+            if not any(cleaned):
+                continue
+            if not header:
+                header = cleaned
+            else:
+                rows.append(cleaned)
+    return header, rows
+
+
+def header_match_score(gt_header: List[str], pred_header: List[str]) -> float:
+    if not gt_header or not pred_header:
+        return 0.0
+    normalized_gt = [cell.strip().lower() for cell in gt_header]
+    normalized_pred = [cell.strip().lower() for cell in pred_header]
+    return 1.0 if normalized_gt == normalized_pred else 0.0
+
+
+def _to_float(cell: str) -> Optional[float]:
+    cell = cell.strip()
+    if not cell:
+        return None
+    normalized = cell.replace(",", "")
+    if normalized.endswith("%"):
+        normalized = normalized[:-1]
+    if normalized.startswith("$"):
+        normalized = normalized[1:]
+    try:
+        return float(normalized)
+    except ValueError:
+        return None
+
+
+def _build_row_map(rows: List[List[str]]) -> Dict[str, List[List[str]]]:
+    mapping: Dict[str, List[List[str]]] = {}
+    for row in rows:
+        if not row:
+            continue
+        key = row[0].strip().lower()
+        mapping.setdefault(key, []).append(row[1:])
+    return mapping
+
+
+def _score_row(gt_values: List[str], pred_values: List[str]) -> float:
+    if not gt_values:
+        return 1.0
+
+    per_column_scores: List[float] = []
+    for idx, gt_cell in enumerate(gt_values):
+        pred_cell = pred_values[idx] if idx < len(pred_values) else ""
+        gt_num = _to_float(gt_cell)
+        pred_num = _to_float(pred_cell)
+        if gt_num is None or pred_num is None:
+            per_column_scores.append(0.0)
+            continue
+        denom = abs(gt_num) + abs(pred_num)
+        if denom == 0:
+            per_column_scores.append(1.0)
+            continue
+        smape = 2.0 * abs(pred_num - gt_num) / denom
+        per_column_scores.append(max(0.0, 1.0 - smape / 2.0))
+
+    if not per_column_scores:
+        return 0.0
+    return sum(per_column_scores) / len(per_column_scores)
+
+
+def visualize_difference(
+    gt_csv_path: Path,
+    pred_csv_path: Path,
+    example_id: str,
+    output_dir: Path,
+    ignore_header_mismatch: bool = False,
+    verbose: bool = False,
+) -> Optional[Path]:
+    def _viz_skip(reason: str) -> None:
+        if verbose:
+            print(f"[viz] skip {example_id}: {reason}", file=sys.stderr)
+
+    if plt is None:
+        _viz_skip("matplotlib not available")
+        return None
+
+    try:
+        gt_header, gt_rows = read_csv_table(gt_csv_path)
+        pred_header, pred_rows = read_csv_table(pred_csv_path)
+    except Exception:
+        _viz_skip("failed to read CSVs")
+        return None
+
+    if not gt_header or not pred_header:
+        _viz_skip("missing headers")
+        return None
+
+    if header_match_score(gt_header, pred_header) < 1.0 and not ignore_header_mismatch:
+        _viz_skip("header mismatch (use --visualize-allow-header-mismatch to override)")
+        return None
+
+    if len(gt_header) <= 1:
+        _viz_skip("no data columns in GT header")
+        return None
+
+    columns = gt_header[1:]
+    gt_series: Dict[str, List[float]] = {col: [] for col in columns}
+    pred_series: Dict[str, List[float]] = {col: [] for col in columns}
+    diff_series: Dict[str, List[float]] = {col: [] for col in columns}
+    x_labels: List[str] = []
+
+    pred_map = _build_row_map(pred_rows)
+    pred_consumed: Dict[str, int] = {}
+
+    for gt_row in gt_rows:
+        if not gt_row:
+            continue
+        x_label = gt_row[0].strip()
+        key = x_label.lower()
+        pred_entries = pred_map.get(key, [])
+        pred_idx = pred_consumed.get(key, 0)
+        pred_values = pred_entries[pred_idx] if pred_idx < len(pred_entries) else []
+        pred_consumed[key] = pred_idx + 1
+
+        x_labels.append(x_label or f"row_{len(x_labels) + 1}")
+
+        for col_idx, col_name in enumerate(columns):
+            gt_val = _to_float(gt_row[col_idx + 1]) if col_idx + 1 < len(gt_row) else None
+            pred_val = _to_float(pred_values[col_idx]) if col_idx < len(pred_values) else None
+            gt_float = gt_val if gt_val is not None else math.nan
+            pred_float = pred_val if pred_val is not None else math.nan
+
+            if math.isnan(gt_float):
+                # If GT is missing, treat as zero difference but keep nan for plotting gaps
+                diff = math.nan
+            elif math.isnan(pred_float):
+                diff = math.nan
+            else:
+                diff = pred_float - gt_float
+
+            gt_series[col_name].append(gt_float)
+            pred_series[col_name].append(pred_float)
+            diff_series[col_name].append(diff)
+
+    if not x_labels:
+        _viz_skip("no x labels in GT rows")
+        return None
+
+    num_series = len(columns)
+    if num_series == 0:
+        _viz_skip("no numeric series to plot")
+        return None
+
+    x_positions = list(range(len(x_labels)))
+    fig_height = max(3.0, 3.0 * num_series)
+    fig, axes = plt.subplots(num_series, 1, sharex=True, figsize=(11, fig_height))
+    if num_series == 1:
+        axes = [axes]
+
+    for ax, col_name in zip(axes, columns):
+        gt_values = gt_series[col_name]
+        pred_values = pred_series[col_name]
+        diff_values = diff_series[col_name]
+
+        ax.plot(x_positions, gt_values, marker="o", linewidth=1.5, label="Ground Truth")
+        ax.plot(x_positions, pred_values, marker="o", linewidth=1.5, label="Prediction")
+        ax.set_ylabel(col_name)
+        ax.grid(True, axis="y", alpha=0.3)
+
+        has_diff = any(not math.isnan(v) for v in diff_values)
+        legend_handles, legend_labels = ax.get_legend_handles_labels()
+
+        if has_diff:
+            ax2 = ax.twinx()
+            ax2.plot(x_positions, diff_values, linestyle="--", color="tab:red", marker="x", linewidth=1.2, label="Pred - GT")
+            ax2.axhline(0.0, color="tab:red", linewidth=0.8, alpha=0.4)
+            ax2.set_ylabel("Pred - GT")
+            handles2, labels2 = ax2.get_legend_handles_labels()
+            legend_handles += handles2
+            legend_labels += labels2
+
+        if legend_handles:
+            ax.legend(legend_handles, legend_labels, loc="upper left")
+
+    axes[-1].set_xticks(x_positions)
+    axes[-1].set_xticklabels(x_labels, rotation=45, ha="right")
+    axes[-1].set_xlabel(gt_header[0])
+
+    fig.suptitle(f"{example_id}: Ground Truth vs Prediction", fontsize=14)
+    fig.tight_layout(rect=(0, 0, 1, 0.96))
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    out_path = output_dir / f"{example_id}.png"
+    fig.savefig(out_path, dpi=150)
+    plt.close(fig)
+    return out_path
+
+
+def evaluate_predictions(gt_csv_path: Path, pred_csv_path: Path) -> float:
+    gt_header, gt_rows = read_csv_table(gt_csv_path)
+    pred_header, pred_rows = read_csv_table(pred_csv_path)
+    if not gt_header or not pred_header:
+        return 0.0
+
+    header_score = header_match_score(gt_header, pred_header)
+
+    gt_map = _build_row_map(gt_rows)
+    pred_map = _build_row_map(pred_rows)
+
+    row_scores: List[float] = []
+    for key, gt_entries in gt_map.items():
+        pred_entries = pred_map.get(key, [])
+        for idx, gt_values in enumerate(gt_entries):
+            pred_values = pred_entries[idx] if idx < len(pred_entries) else []
+            row_scores.append(_score_row(gt_values, pred_values))
+
+    content_score = sum(row_scores) / len(row_scores) if row_scores else 0.0
+    return 0.2 * header_score + 0.8 * content_score
+
+
+def process_one(
+    extractor: VLMExtractor, base_dir: Path, example_id: str, image_rel: Path, gt_table_rel: Path, output_dir: Path
+) -> Tuple[str, float, Path, Path]:
+    image_path = base_dir / image_rel
+    gt_csv_path = base_dir / gt_table_rel
+    pred_csv_text = extractor.image_to_csv(image_path)
+    pred_path = write_csv(output_dir, example_id, pred_csv_text)
+    score = evaluate_predictions(gt_csv_path, pred_path)
+    return example_id, score, pred_path, gt_csv_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Evaluate VLM extraction: image -> CSV")
+    parser.add_argument("--data-dir", type=str, default="subset_line_100")
+    parser.add_argument("--index", type=str, default="index.csv")
+    parser.add_argument("--out-dir", type=str, default="predictions")
+    parser.add_argument("--max-samples", type=int, default=100)
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument(
+        "--visualize-dir",
+        type=str,
+        default=None,
+        help="Directory where GT vs prediction plots will be saved (requires matplotlib).",
+    )
+    parser.add_argument(
+        "--visualize-max",
+        type=int,
+        default=1,
+        help="Maximum number of plots to generate when --visualize-dir is set. Use 0 for no limit.",
+    )
+    parser.add_argument(
+        "--visualize-allow-header-mismatch",
+        action="store_true",
+        help="If set, still plot GT vs prediction even when headers differ.",
+    )
+    parser.add_argument("--visualize-verbose", action="store_true", help="Print reasons when a visualization is skipped.")
+    args = parser.parse_args()
+
+    if not os.getenv("OPENAI_API_KEY"):
+        print("[error] OPENAI_API_KEY not set in environment", file=sys.stderr)
+        sys.exit(1)
+
+    base_dir = Path(args.data_dir)
+    index_path = base_dir / args.index
+    if not index_path.exists():
+        print(f"[error] index.csv not found at {index_path}", file=sys.stderr)
+        sys.exit(1)
+
+    rows = read_index(index_path)[: args.max_samples]
+    extractor = VLMExtractor()
+
+    visualize_dir: Optional[Path] = Path(args.visualize_dir) if args.visualize_dir else None
+    visualize_max = max(0, args.visualize_max)
+    if visualize_dir and plt is None:
+        print("[warn] matplotlib not available; skipping visualization.", file=sys.stderr)
+        visualize_dir = None
+
+    print(f"[setup] evaluating {len(rows)} samples using {extractor.model} …", flush=True)
+    start = time.time()
+    scores: List[float] = []
+    saved_visualizations = 0
+
+    with ThreadPoolExecutor(max_workers=max(1, args.num_workers)) as pool:
+        futures = [
+            pool.submit(process_one, extractor, base_dir, example_id, image_rel, gt_table_rel, Path(args.out_dir))
+            for (example_id, image_rel, gt_table_rel) in rows
+        ]
+
+        try:
+            for idx, fut in enumerate(as_completed(futures), 1):
+                try:
+                    example_id, score, pred_path, gt_csv_path = fut.result()
+                    scores.append(score)
+                    if visualize_dir and (visualize_max == 0 or saved_visualizations < visualize_max):
+                        out_path = visualize_difference(
+                            gt_csv_path,
+                            pred_path,
+                            example_id,
+                            visualize_dir,
+                            ignore_header_mismatch=args.visualize_allow_header_mismatch,
+                            verbose=args.visualize_verbose,
+                        )
+                        if out_path is not None:
+                            saved_visualizations += 1
+                            print(f"[viz] saved {out_path}", flush=True)
+                    if idx % 5 == 0 or idx == len(rows):
+                        elapsed = time.time() - start
+                        avg = sum(scores) / len(scores) if scores else 0.0
+                        print(f"[progress] {idx}/{len(rows)} done, avg score: {avg:.4f}, elapsed {elapsed:.1f}s", flush=True)
+                except Exception as e:
+                    print(f"[error] failed on sample {idx}: {e}", file=sys.stderr)
+        except KeyboardInterrupt:
+            print("\n[warn] interrupted by user", file=sys.stderr)
+            sys.exit(1)
+
+    final_score = sum(scores) / len(scores) if scores else 0.0
+
+    # Apply cost cap: accuracy is zeroed if average cost/query exceeds $0.02
+    avg_cost_per_query = (extractor.total_cost_usd / extractor.num_queries) if getattr(extractor, "num_queries", 0) else 0.0
+    if avg_cost_per_query > 0.02:
+        print(f"[cost] avg ${avg_cost_per_query:.4f}/query exceeds $0.02 cap; accuracy set to 0.0", flush=True)
+        final_score = 0.0
+    else:
+        print(f"[cost] avg ${avg_cost_per_query:.4f}/query within cap", flush=True)
+
+    print(f"accuracy: {final_score:.4f}")
+
+
+if __name__ == "__main__":
+    main()

weco-0.3.2/examples/extract-line-plot/guide.md

@@ -0,0 +1,53 @@
+# Constraints
+- Make sure the cost tracking is correct
+- The average cost per query should be lower than $0.02 per query, the baseline has around $0.007 per query
+
+# Here's a list of ideas to be tried:
+1. Try more recent models like gpt-5, note gpt-5 don't have temperature argument
+2. Try response API as below:
+
+<api>
+Responses API (concise)
+
+Endpoint:
+- POST https://api.openai.com/v1/responses
+
+Purpose:
+- Generate text or JSON from text and/or image inputs. Supports optional tools.
+
+Minimum request fields:
+- model: string (e.g., gpt-5, gpt-4o)
+- input: string | array (text and/or image items)
+
+Useful options:
+- max_output_tokens: integer
+- temperature or top_p: sampling controls (non-reasoning models)
+- reasoning: object (gpt-5, o-series) — controls for reasoning models
+- tools: array (web_search, file_search, function calls)
+- tool_choice: string | object
+- stream: boolean (default false)
+- store: boolean (default true)
+- parallel_tool_calls: boolean (default true)
+
+Conversation (optional):
+- conversation: string | object, or previous_response_id: string
+
+Notes:
+- gpt-5 does not support temperature; prefer reasoning options instead.
+
+Example request:
+
+```json
+{
+  "model": "gpt-5",
+  "input": "Explain how to extract a line plot from an image.",
+  "max_output_tokens": 400,
+  "reasoning": { "effort": "medium" },
+  "stream": false
+}
+```
+</api>
+
+3. Try playing with paramters like reasoning efforts
+4. Try to build tools for response api to use, or allow python interpretor
+5. Try to add preprocessing or post processing pipelines