runner.py

"""Batch similarity experiment. Usage: OPENROUTER_API_KEY=... python runner.py"""
import argparse
import csv
import json
import os
from pathlib import Path
from embedding_similarity import similarity

ROOT = Path(__file__).resolve().parent


def run(cases_path: Path, output_path: Path, token: str, limit: int | None = None):
    cases = json.loads(cases_path.read_text(encoding="utf-8"))
    rows = []
    cache = {}
    for item in cases[:limit]:
        a, b = item["left"], item["right"]
        texts = []
        for filename in (a, b):
            if filename not in cache:
                cache[filename] = (ROOT / "samples" / filename).read_text(encoding="utf-8")
            texts.append(cache[filename])
        try:
            score = similarity(texts[0], texts[1], token)
            error = ""
        except Exception as exc:
            score = ""
            error = str(exc)
        row = {**item, "similarity": score, "error": error}
        rows.append(row)
        print(f'{item["id"]} {item["expected"]:8s} score={str(score):>10s}  {a} <> {b}' + (f' ERROR: {error}' if error else ''))
        if error:
            print("Stopping on API error to avoid repeated requests.")
            break
    with output_path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["id", "left", "right", "expected", "category", "notes", "similarity", "error"])
        writer.writeheader()
        writer.writerows(rows)
    print(f"Saved {len(rows)} results: {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=None, help="Run only first N pairs to minimize costs")
    parser.add_argument("--cases", type=Path, default=ROOT / "cases.json")
    parser.add_argument("--output", type=Path, default=ROOT / "results.csv")
    args = parser.parse_args()
    key = os.environ.get("OPENROUTER_API_KEY", "")
    if not key:
        parser.error("Set OPENROUTER_API_KEY environment variable; never put keys in command line args")
    run(args.cases, args.output, key, args.limit)