wehub-resource-sync
680 lines
21 KiB
Python
680 lines
21 KiB
Python
#!/usr/bin/env python3
|
|
"""Clean query-only training rows toward a SimpleQA-like style.
|
|
|
|
This script reads one or more JSONL files that contain `query` fields and uses
|
|
Gemini to judge each query on two axes:
|
|
1. Naturalness: does it sound like a real user question?
|
|
2. SimpleQA style fit: does it resemble the style of SimpleQA factoid prompts?
|
|
|
|
The model only sees the query text. It does not inspect images or answers.
|
|
|
|
Outputs:
|
|
- cleaned JSONL with original rows preserved
|
|
- review JSONL with model scores / keep decisions
|
|
- summary JSON with before/after stats and token usage
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
import random
|
|
import re
|
|
import threading
|
|
import time
|
|
from collections import Counter
|
|
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
|
from pathlib import Path
|
|
|
|
|
|
DEFAULT_INPUT_GLOB = (
|
|
"training/data/lite-query-v2-full-filtered-hn-v2-chunks/chunk_*/filtered_hn.jsonl"
|
|
)
|
|
DEFAULT_SIMPLEQA_PATH = (
|
|
"/home/user/wiki-screenshot/eval/simpleqa_query_image_pairs.json"
|
|
)
|
|
|
|
MODEL_PRICING = {
|
|
"gemini-2.0-flash-001": {"input_per_m": 0.10, "output_per_m": 0.40},
|
|
"gemini-2.0-flash": {"input_per_m": 0.10, "output_per_m": 0.40},
|
|
"gemini-3.1-flash": {"input_per_m": 0.25, "output_per_m": 1.50},
|
|
}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--input-glob",
|
|
default=DEFAULT_INPUT_GLOB,
|
|
help="Glob pattern for input JSONL files.",
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
required=True,
|
|
help="Output JSONL path for cleaned rows.",
|
|
)
|
|
parser.add_argument(
|
|
"--reviews-output",
|
|
default=None,
|
|
help="Optional JSONL path for model review decisions.",
|
|
)
|
|
parser.add_argument(
|
|
"--summary-output",
|
|
default=None,
|
|
help="Optional JSON path for summary stats.",
|
|
)
|
|
parser.add_argument(
|
|
"--simpleqa-path",
|
|
default=DEFAULT_SIMPLEQA_PATH,
|
|
help="Path to SimpleQA pair JSON used for style references.",
|
|
)
|
|
parser.add_argument(
|
|
"--target-count",
|
|
type=int,
|
|
default=50000,
|
|
help="Approximate number of rows to keep.",
|
|
)
|
|
parser.add_argument(
|
|
"--batch-size",
|
|
type=int,
|
|
default=20,
|
|
help="Queries per Gemini request.",
|
|
)
|
|
parser.add_argument(
|
|
"--concurrency",
|
|
type=int,
|
|
default=8,
|
|
help="Concurrent Gemini requests.",
|
|
)
|
|
parser.add_argument(
|
|
"--few-shot-count",
|
|
type=int,
|
|
default=12,
|
|
help="Number of SimpleQA examples shown in the prompt.",
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
default="gemini-2.0-flash-001",
|
|
help="Gemini model name. Use an accessible model in your project.",
|
|
)
|
|
parser.add_argument("--gemini-project", default="wise-coyote-478119-h0")
|
|
parser.add_argument("--gemini-location", default="global")
|
|
parser.add_argument("--max-retries", type=int, default=5)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument(
|
|
"--min-naturalness",
|
|
type=int,
|
|
default=4,
|
|
help="Minimum Gemini naturalness score for direct keep.",
|
|
)
|
|
parser.add_argument(
|
|
"--min-style-fit",
|
|
type=int,
|
|
default=4,
|
|
help="Minimum Gemini SimpleQA style score for direct keep.",
|
|
)
|
|
parser.add_argument(
|
|
"--dedupe-query",
|
|
action="store_true",
|
|
help="Keep only the highest-scoring row per normalized query.",
|
|
)
|
|
parser.add_argument(
|
|
"--resume",
|
|
action="store_true",
|
|
help="Reuse existing review file and skip already-scored row ids.",
|
|
)
|
|
parser.add_argument(
|
|
"--limit",
|
|
type=int,
|
|
default=0,
|
|
help="Optional cap on total input rows for smoke tests.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def init_usage() -> dict:
|
|
return {
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"calls": 0,
|
|
}
|
|
|
|
|
|
def build_client(args: argparse.Namespace) -> dict:
|
|
os.environ.setdefault("GOOGLE_CLOUD_PROJECT", args.gemini_project)
|
|
os.environ.setdefault("GOOGLE_GENAI_USE_VERTEXAI", "true")
|
|
os.environ.setdefault("GOOGLE_CLOUD_LOCATION", args.gemini_location)
|
|
|
|
from google import genai
|
|
from google.genai.types import HttpOptions
|
|
|
|
client = genai.Client(http_options=HttpOptions(api_version="v1"))
|
|
return {
|
|
"client": client,
|
|
"usage": init_usage(),
|
|
"usage_lock": threading.Lock(),
|
|
}
|
|
|
|
|
|
def update_usage(
|
|
client_ctx: dict, prompt_tokens: int = 0, completion_tokens: int = 0
|
|
) -> None:
|
|
with client_ctx["usage_lock"]:
|
|
client_ctx["usage"]["prompt_tokens"] += int(prompt_tokens or 0)
|
|
client_ctx["usage"]["completion_tokens"] += int(completion_tokens or 0)
|
|
client_ctx["usage"]["calls"] += 1
|
|
|
|
|
|
def parse_json_from_text(text: str):
|
|
text = text.strip()
|
|
if not text:
|
|
raise ValueError("empty response")
|
|
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
match = re.search(r"(\[[\s\S]*\]|\{[\s\S]*\})", text)
|
|
if not match:
|
|
raise ValueError("no JSON object or array found")
|
|
return json.loads(match.group(1))
|
|
|
|
|
|
def call_gemini_json(client_ctx: dict, model: str, prompt: str, max_retries: int):
|
|
from google.genai.types import GenerateContentConfig
|
|
|
|
config = GenerateContentConfig(
|
|
temperature=0,
|
|
max_output_tokens=8192,
|
|
response_mime_type="application/json",
|
|
)
|
|
for attempt in range(1, max_retries + 1):
|
|
try:
|
|
resp = client_ctx["client"].models.generate_content(
|
|
model=model,
|
|
contents=prompt,
|
|
config=config,
|
|
)
|
|
usage = getattr(resp, "usage_metadata", None)
|
|
if usage is not None:
|
|
update_usage(
|
|
client_ctx,
|
|
prompt_tokens=getattr(usage, "prompt_token_count", 0),
|
|
completion_tokens=getattr(usage, "candidates_token_count", 0),
|
|
)
|
|
text = getattr(resp, "text", "") or ""
|
|
return parse_json_from_text(text)
|
|
except Exception:
|
|
if attempt == max_retries:
|
|
raise
|
|
time.sleep(min(2**attempt, 20))
|
|
raise RuntimeError("unreachable")
|
|
|
|
|
|
def iter_rows(path: Path):
|
|
with path.open() as f:
|
|
for line_no, line in enumerate(f, 1):
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
yield line_no, json.loads(line)
|
|
|
|
|
|
def load_input_rows(args: argparse.Namespace) -> list[dict]:
|
|
rows = []
|
|
matched = [Path(path) for path in sorted(glob.glob(args.input_glob))]
|
|
if not matched:
|
|
raise FileNotFoundError(f"No files matched --input-glob={args.input_glob}")
|
|
|
|
row_id = 0
|
|
for path in matched:
|
|
for line_no, payload in iter_rows(path):
|
|
query = payload.get("query") or payload.get("question")
|
|
if not isinstance(query, str) or not query.strip():
|
|
continue
|
|
rows.append(
|
|
{
|
|
"row_id": row_id,
|
|
"source_file": str(path),
|
|
"source_line": line_no,
|
|
"query": query.strip(),
|
|
"payload": payload,
|
|
}
|
|
)
|
|
row_id += 1
|
|
if args.limit > 0 and len(rows) >= args.limit:
|
|
return rows
|
|
return rows
|
|
|
|
|
|
def normalize_query(query: str) -> str:
|
|
return " ".join(query.lower().split())
|
|
|
|
|
|
def question_start_bucket(query: str) -> str:
|
|
q = query.strip().lower()
|
|
for prefix in [
|
|
"what",
|
|
"who",
|
|
"which",
|
|
"when",
|
|
"where",
|
|
"why",
|
|
"how",
|
|
"in which",
|
|
"in what",
|
|
"on what",
|
|
"what is",
|
|
"what was",
|
|
]:
|
|
if q.startswith(prefix):
|
|
return prefix
|
|
parts = re.findall(r"[a-z0-9']+", q)
|
|
return parts[0] if parts else "<empty>"
|
|
|
|
|
|
def load_simpleqa_references(path: Path, few_shot_count: int, seed: int) -> list[dict]:
|
|
with path.open() as f:
|
|
data = json.load(f)
|
|
|
|
by_bucket = {}
|
|
for item in data:
|
|
question = item.get("question")
|
|
answer = item.get("answer")
|
|
if not isinstance(question, str) or not question.strip():
|
|
continue
|
|
bucket = question_start_bucket(question)
|
|
by_bucket.setdefault(bucket, []).append(
|
|
{
|
|
"question": question.strip(),
|
|
"answer": (answer or "").strip(),
|
|
"topic": item.get("topic") or "Unknown",
|
|
}
|
|
)
|
|
|
|
rng = random.Random(seed)
|
|
for values in by_bucket.values():
|
|
rng.shuffle(values)
|
|
|
|
ordered_buckets = sorted(by_bucket, key=lambda key: (-len(by_bucket[key]), key))
|
|
refs = []
|
|
while len(refs) < few_shot_count and ordered_buckets:
|
|
next_round = []
|
|
for bucket in ordered_buckets:
|
|
values = by_bucket[bucket]
|
|
if values:
|
|
refs.append(values.pop())
|
|
if values:
|
|
next_round.append(bucket)
|
|
if len(refs) >= few_shot_count:
|
|
break
|
|
ordered_buckets = next_round
|
|
return refs
|
|
|
|
|
|
def build_prompt(reference_examples: list[dict], batch: list[dict]) -> str:
|
|
ref_lines = []
|
|
for idx, ref in enumerate(reference_examples, 1):
|
|
ref_lines.append(
|
|
f"{idx}. topic={ref['topic']}\n"
|
|
f" question={json.dumps(ref['question'], ensure_ascii=False)}\n"
|
|
f" answer={json.dumps(ref['answer'], ensure_ascii=False)}"
|
|
)
|
|
|
|
candidate_json = json.dumps(
|
|
[{"id": row["row_id"], "query": row["query"]} for row in batch],
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
|
|
return f"""You are cleaning a synthetic screenshot-retrieval training set.
|
|
|
|
Goal: keep only queries that read naturally and resemble SimpleQA-style factoid questions.
|
|
Judge ONLY the query text. Do not assume access to images, answers, or metadata.
|
|
|
|
High-quality queries usually:
|
|
- sound like something a real user would ask
|
|
- are single-hop factoid questions with short answers
|
|
- include natural disambiguating context when useful (time, role, location, quoted title, relation)
|
|
- feel similar to the reference questions below
|
|
|
|
Reject queries that are:
|
|
- awkward, templatic, or annotation-like
|
|
- unnaturally phrased, especially stiff starts like "In what..." or "At which..." when a plain wording would be more natural
|
|
- keywordy, malformed, or obviously synthetic
|
|
- broad explanatory prompts, opinion questions, yes/no questions, or multi-hop questions
|
|
|
|
Reference SimpleQA-style examples:
|
|
{chr(10).join(ref_lines)}
|
|
|
|
Now score these candidate queries.
|
|
|
|
Return ONLY a JSON array. One object per candidate with exactly these keys:
|
|
- id: integer
|
|
- naturalness: integer 1-5
|
|
- simpleqa_style_fit: integer 1-5
|
|
- keep: boolean
|
|
- reason: short string (<=12 words)
|
|
|
|
Scoring guide:
|
|
- 5 = excellent
|
|
- 4 = good / clearly keepable
|
|
- 3 = borderline
|
|
- 2 = weak
|
|
- 1 = poor
|
|
|
|
Candidates:
|
|
{candidate_json}
|
|
"""
|
|
|
|
|
|
def sanitize_decision(raw: dict, row_id: int) -> dict:
|
|
naturalness = raw.get("naturalness", 0)
|
|
style_fit = raw.get("simpleqa_style_fit", 0)
|
|
keep = raw.get("keep", False)
|
|
reason = raw.get("reason", "")
|
|
|
|
try:
|
|
naturalness = int(naturalness)
|
|
except Exception:
|
|
naturalness = 0
|
|
try:
|
|
style_fit = int(style_fit)
|
|
except Exception:
|
|
style_fit = 0
|
|
if isinstance(keep, str):
|
|
keep = keep.strip().lower() in {"true", "yes", "1", "keep"}
|
|
|
|
return {
|
|
"id": int(row_id),
|
|
"naturalness": max(0, min(5, naturalness)),
|
|
"simpleqa_style_fit": max(0, min(5, style_fit)),
|
|
"keep": bool(keep),
|
|
"reason": str(reason).strip()[:200],
|
|
}
|
|
|
|
|
|
def score_batch(
|
|
client_ctx: dict,
|
|
args: argparse.Namespace,
|
|
references: list[dict],
|
|
batch: list[dict],
|
|
) -> list[dict]:
|
|
prompt = build_prompt(references, batch)
|
|
parsed = call_gemini_json(client_ctx, args.model, prompt, args.max_retries)
|
|
if not isinstance(parsed, list):
|
|
raise ValueError("Gemini response is not a JSON array")
|
|
|
|
by_id = {}
|
|
for item in parsed:
|
|
if not isinstance(item, dict) or "id" not in item:
|
|
continue
|
|
by_id[int(item["id"])] = sanitize_decision(item, item["id"])
|
|
|
|
decisions = []
|
|
for row in batch:
|
|
decision = by_id.get(row["row_id"])
|
|
if decision is None:
|
|
decision = {
|
|
"id": row["row_id"],
|
|
"naturalness": 0,
|
|
"simpleqa_style_fit": 0,
|
|
"keep": False,
|
|
"reason": "missing_from_model_output",
|
|
}
|
|
decisions.append(decision)
|
|
return decisions
|
|
|
|
|
|
def load_existing_reviews(path: Path) -> dict[int, dict]:
|
|
existing = {}
|
|
if not path.exists():
|
|
return existing
|
|
with path.open() as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
item = json.loads(line)
|
|
existing[int(item["row_id"])] = item
|
|
return existing
|
|
|
|
|
|
def review_rows(
|
|
rows: list[dict],
|
|
args: argparse.Namespace,
|
|
client_ctx: dict,
|
|
references: list[dict],
|
|
reviews_path: Path,
|
|
) -> dict[int, dict]:
|
|
existing = load_existing_reviews(reviews_path) if args.resume else {}
|
|
pending = [row for row in rows if row["row_id"] not in existing]
|
|
total = len(rows)
|
|
|
|
if pending:
|
|
reviews_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with reviews_path.open("a") as reviews_file:
|
|
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
|
|
futures = {}
|
|
pending_batches = [
|
|
pending[start : start + args.batch_size]
|
|
for start in range(0, len(pending), args.batch_size)
|
|
]
|
|
next_batch_idx = 0
|
|
|
|
while (
|
|
next_batch_idx < len(pending_batches)
|
|
and len(futures) < args.concurrency
|
|
):
|
|
batch = pending_batches[next_batch_idx]
|
|
future = executor.submit(
|
|
score_batch, client_ctx, args, references, batch
|
|
)
|
|
futures[future] = batch
|
|
next_batch_idx += 1
|
|
|
|
completed = 0
|
|
while futures:
|
|
done, _ = wait(futures.keys(), return_when=FIRST_COMPLETED)
|
|
for future in done:
|
|
batch = futures.pop(future)
|
|
decisions = future.result()
|
|
for row, decision in zip(batch, decisions):
|
|
review = {
|
|
"row_id": row["row_id"],
|
|
"query": row["query"],
|
|
"source_file": row["source_file"],
|
|
"source_line": row["source_line"],
|
|
**decision,
|
|
}
|
|
existing[row["row_id"]] = review
|
|
reviews_file.write(
|
|
json.dumps(review, ensure_ascii=False) + "\n"
|
|
)
|
|
reviews_file.flush()
|
|
completed += len(batch)
|
|
print(
|
|
f"Reviewed {len(existing)}/{total} rows "
|
|
f"(new {completed}/{len(pending)})",
|
|
flush=True,
|
|
)
|
|
if next_batch_idx < len(pending_batches):
|
|
next_batch = pending_batches[next_batch_idx]
|
|
next_future = executor.submit(
|
|
score_batch, client_ctx, args, references, next_batch
|
|
)
|
|
futures[next_future] = next_batch
|
|
next_batch_idx += 1
|
|
return existing
|
|
|
|
|
|
def candidate_priority(review: dict) -> tuple:
|
|
return (
|
|
int(review["keep"]),
|
|
int(review["naturalness"]) + int(review["simpleqa_style_fit"]),
|
|
int(review["simpleqa_style_fit"]),
|
|
int(review["naturalness"]),
|
|
-int(review["row_id"]),
|
|
)
|
|
|
|
|
|
def select_rows(
|
|
rows: list[dict], reviews: dict[int, dict], args: argparse.Namespace
|
|
) -> list[dict]:
|
|
candidates = []
|
|
for row in rows:
|
|
review = reviews.get(row["row_id"])
|
|
if not review:
|
|
continue
|
|
direct_keep = (
|
|
review["keep"]
|
|
and review["naturalness"] >= args.min_naturalness
|
|
and review["simpleqa_style_fit"] >= args.min_style_fit
|
|
)
|
|
if direct_keep:
|
|
candidates.append((row, review))
|
|
|
|
if args.dedupe_query:
|
|
by_query = {}
|
|
for row, review in candidates:
|
|
key = normalize_query(row["query"])
|
|
current = by_query.get(key)
|
|
if current is None or candidate_priority(review) > candidate_priority(
|
|
current[1]
|
|
):
|
|
by_query[key] = (row, review)
|
|
candidates = list(by_query.values())
|
|
|
|
candidates.sort(key=lambda item: candidate_priority(item[1]), reverse=True)
|
|
|
|
if args.target_count > 0 and len(candidates) > args.target_count:
|
|
candidates = candidates[: args.target_count]
|
|
|
|
return [row for row, _ in candidates]
|
|
|
|
|
|
def word_count(query: str) -> int:
|
|
return len(re.findall(r"[A-Za-z0-9']+", query))
|
|
|
|
|
|
def compute_query_stats(queries: list[str]) -> dict:
|
|
if not queries:
|
|
return {
|
|
"count": 0,
|
|
"avg_words": 0.0,
|
|
"avg_chars": 0.0,
|
|
"top_starts": [],
|
|
"has_quote_pct": 0.0,
|
|
"has_year_pct": 0.0,
|
|
}
|
|
|
|
starts = Counter(question_start_bucket(query) for query in queries)
|
|
word_counts = [word_count(query) for query in queries]
|
|
char_counts = [len(query) for query in queries]
|
|
has_quote = sum('"' in query or "'" in query for query in queries)
|
|
has_year = sum(
|
|
bool(re.search(r"\b(1[0-9]{3}|20[0-2][0-9])\b", query)) for query in queries
|
|
)
|
|
return {
|
|
"count": len(queries),
|
|
"avg_words": round(sum(word_counts) / len(word_counts), 2),
|
|
"avg_chars": round(sum(char_counts) / len(char_counts), 2),
|
|
"top_starts": starts.most_common(12),
|
|
"has_quote_pct": round(100 * has_quote / len(queries), 2),
|
|
"has_year_pct": round(100 * has_year / len(queries), 2),
|
|
}
|
|
|
|
|
|
def estimated_cost_usd(model: str, usage: dict) -> float | None:
|
|
pricing = MODEL_PRICING.get(model)
|
|
if not pricing:
|
|
return None
|
|
in_cost = usage["prompt_tokens"] / 1_000_000 * pricing["input_per_m"]
|
|
out_cost = usage["completion_tokens"] / 1_000_000 * pricing["output_per_m"]
|
|
return round(in_cost + out_cost, 6)
|
|
|
|
|
|
def write_jsonl(path: Path, rows: list[dict]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row["payload"], ensure_ascii=False) + "\n")
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
output_path = Path(args.output)
|
|
reviews_path = (
|
|
Path(args.reviews_output)
|
|
if args.reviews_output
|
|
else output_path.with_suffix(".reviews.jsonl")
|
|
)
|
|
summary_path = (
|
|
Path(args.summary_output)
|
|
if args.summary_output
|
|
else output_path.with_suffix(".summary.json")
|
|
)
|
|
|
|
random.seed(args.seed)
|
|
|
|
rows = load_input_rows(args)
|
|
print(f"Loaded {len(rows)} rows from {args.input_glob}", flush=True)
|
|
|
|
references = load_simpleqa_references(
|
|
Path(args.simpleqa_path), args.few_shot_count, args.seed
|
|
)
|
|
print(f"Loaded {len(references)} SimpleQA reference examples", flush=True)
|
|
|
|
client_ctx = build_client(args)
|
|
reviews = review_rows(rows, args, client_ctx, references, reviews_path)
|
|
selected = select_rows(rows, reviews, args)
|
|
write_jsonl(output_path, selected)
|
|
|
|
input_queries = [row["query"] for row in rows]
|
|
output_queries = [row["query"] for row in selected]
|
|
summary = {
|
|
"model": args.model,
|
|
"input_glob": args.input_glob,
|
|
"simpleqa_path": args.simpleqa_path,
|
|
"total_input_rows": len(rows),
|
|
"reviewed_rows": len(reviews),
|
|
"selected_rows": len(selected),
|
|
"target_count": args.target_count,
|
|
"dedupe_query": args.dedupe_query,
|
|
"min_naturalness": args.min_naturalness,
|
|
"min_style_fit": args.min_style_fit,
|
|
"batch_size": args.batch_size,
|
|
"concurrency": args.concurrency,
|
|
"query_stats_before": compute_query_stats(input_queries),
|
|
"query_stats_after": compute_query_stats(output_queries),
|
|
"usage": {
|
|
**client_ctx["usage"],
|
|
"estimated_cost_usd": estimated_cost_usd(args.model, client_ctx["usage"]),
|
|
},
|
|
}
|
|
summary_path.parent.mkdir(parents=True, exist_ok=True)
|
|
summary_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n")
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"output": str(output_path),
|
|
"reviews_output": str(reviews_path),
|
|
"summary_output": str(summary_path),
|
|
"selected_rows": len(selected),
|
|
"estimated_cost_usd": summary["usage"]["estimated_cost_usd"],
|
|
},
|
|
indent=2,
|
|
ensure_ascii=False,
|
|
),
|
|
flush=True,
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|