init
This commit is contained in:
tokuzou0829
@@ -0,0 +1,757 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Incrementally fetch new ASMR works and update web/data assets."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import mmap
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
ROOT_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(ROOT_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(ROOT_DIR))
|
||||
|
||||
from fetch_asmr_works import fetch_page # noqa: E402
|
||||
from build_vector_data import embedding_text, normalize_work # noqa: E402
|
||||
from build_vector_data_remote import ( # noqa: E402
|
||||
embeddings_url,
|
||||
normalize_vector,
|
||||
request_embeddings_with_retries,
|
||||
)
|
||||
from filter_asmr_works import ( # noqa: E402
|
||||
DEFAULT_EXCLUDE_TAG_IDS,
|
||||
DEFAULT_EXCLUDE_TAGS,
|
||||
filter_works,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Fetch latest ASMR works, filter them, and incrementally update web/data."
|
||||
)
|
||||
parser.add_argument("--input", default="asmr_works.jsonl", help="Local raw ASMR JSONL file.")
|
||||
parser.add_argument(
|
||||
"--filtered-output",
|
||||
default="asmr_works.filtered.jsonl",
|
||||
help="Filtered ASMR JSONL output file.",
|
||||
)
|
||||
parser.add_argument("--data-dir", default="web/data", help="Directory containing web data assets.")
|
||||
parser.add_argument("--order", default="create_date", help="ASMR API order parameter.")
|
||||
parser.add_argument("--sort", default="desc", help="ASMR API sort parameter.")
|
||||
parser.add_argument("--subtitle", default="0", help="ASMR API subtitle parameter.")
|
||||
parser.add_argument("--page-size", type=int, default=100, help="ASMR API pageSize parameter.")
|
||||
parser.add_argument("--timeout", type=float, default=30.0, help="ASMR API request timeout seconds.")
|
||||
parser.add_argument("--retries", type=int, default=3, help="ASMR API retries per failed page.")
|
||||
parser.add_argument(
|
||||
"--max-pages",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Maximum latest pages to scan. By default, stop when a page has no new works.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Rebuild filtered/web data even when no new works are found.",
|
||||
)
|
||||
parser.add_argument("--dry-run", action="store_true", help="Fetch and report without writing files.")
|
||||
parser.add_argument(
|
||||
"--diff-limit",
|
||||
type=int,
|
||||
default=30,
|
||||
help="Maximum rows to print per dry-run diff section. Use 0 for all rows.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--exclude-tag",
|
||||
action="append",
|
||||
default=[],
|
||||
help="Additional tag name to remove during filtering.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--exclude-tag-id",
|
||||
action="append",
|
||||
type=int,
|
||||
default=[],
|
||||
help="Additional tag id to remove during filtering.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--keep-default-exclude-tags",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Keep default 女性向/乙女向 filtering.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-key-env",
|
||||
default="EMBEDDING_API_KEY",
|
||||
help="Environment variable that contains the remote embedding API key.",
|
||||
)
|
||||
parser.add_argument("--embedding-timeout", type=float, default=120.0, help="Embedding request timeout seconds.")
|
||||
parser.add_argument("--embedding-retries", type=int, default=3, help="Embedding retries per failed batch.")
|
||||
parser.add_argument("--embedding-retry-wait", type=float, default=2.0, help="Initial embedding retry wait seconds.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def clean_text(value: Any) -> str:
|
||||
return value.strip() if isinstance(value, str) else ""
|
||||
|
||||
|
||||
def identity_keys(work: dict[str, Any]) -> list[str]:
|
||||
keys: list[str] = []
|
||||
source_id = clean_text(work.get("source_id")) or clean_text(work.get("sourceId"))
|
||||
if source_id:
|
||||
keys.append(f"source:{source_id.lower()}")
|
||||
work_id = work.get("id")
|
||||
if work_id is not None:
|
||||
keys.append(f"id:{work_id}")
|
||||
return keys
|
||||
|
||||
|
||||
def primary_key(work: dict[str, Any]) -> str:
|
||||
keys = identity_keys(work)
|
||||
if not keys:
|
||||
raise RuntimeError(f"work has no id/source_id: {work!r}")
|
||||
return keys[0]
|
||||
|
||||
|
||||
def display_id(work: dict[str, Any]) -> str:
|
||||
return clean_text(work.get("source_id")) or clean_text(work.get("sourceId")) or str(work.get("id") or "unknown")
|
||||
|
||||
|
||||
def display_title(work: dict[str, Any]) -> str:
|
||||
return clean_text(work.get("title")) or "title unknown"
|
||||
|
||||
|
||||
def truncate(value: str, max_length: int = 90) -> str:
|
||||
if len(value) <= max_length:
|
||||
return value
|
||||
return value[: max_length - 3] + "..."
|
||||
|
||||
|
||||
def format_work(work: dict[str, Any]) -> str:
|
||||
date = clean_text(work.get("create_date")) or clean_text(work.get("createDate"))
|
||||
suffix = f" / {date}" if date else ""
|
||||
return f"{display_id(work)} / {truncate(display_title(work))}{suffix}"
|
||||
|
||||
|
||||
def read_jsonl(path: Path) -> list[dict[str, Any]]:
|
||||
works: list[dict[str, Any]] = []
|
||||
with path.open(encoding="utf-8") as file:
|
||||
for line_no, line in enumerate(file, 1):
|
||||
if not line.strip():
|
||||
continue
|
||||
try:
|
||||
value = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise SystemExit(f"{path}:{line_no}: invalid JSON: {exc}") from exc
|
||||
if not isinstance(value, dict):
|
||||
raise SystemExit(f"{path}:{line_no}: expected JSON object")
|
||||
works.append(value)
|
||||
return works
|
||||
|
||||
|
||||
def write_jsonl_atomic(path: Path, works: Iterable[dict[str, Any]]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temp_path = path.with_name(path.name + ".tmp")
|
||||
try:
|
||||
with temp_path.open("w", encoding="utf-8") as file:
|
||||
for work in works:
|
||||
clean_work = {key: value for key, value in work.items() if key != "__line_no"}
|
||||
file.write(json.dumps(clean_work, ensure_ascii=False, separators=(",", ":")))
|
||||
file.write("\n")
|
||||
temp_path.replace(path)
|
||||
except Exception:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def write_json_atomic(path: Path, value: Any, *, indent: int | None = None) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temp_path = path.with_name(path.name + ".tmp")
|
||||
try:
|
||||
with temp_path.open("w", encoding="utf-8") as file:
|
||||
json.dump(value, file, ensure_ascii=False, indent=indent, separators=None if indent else (",", ":"))
|
||||
if indent:
|
||||
file.write("\n")
|
||||
temp_path.replace(path)
|
||||
except Exception:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def fetch_latest_works(args: argparse.Namespace, existing_keys: set[str]) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]], int]:
|
||||
fetch_args = SimpleNamespace(
|
||||
order=args.order,
|
||||
sort=args.sort,
|
||||
subtitle=args.subtitle,
|
||||
page_size=args.page_size,
|
||||
timeout=args.timeout,
|
||||
retries=args.retries,
|
||||
)
|
||||
|
||||
new_works: list[dict[str, Any]] = []
|
||||
fetched_existing_by_key: dict[str, dict[str, Any]] = {}
|
||||
seen_new_keys: set[str] = set()
|
||||
scanned_pages = 0
|
||||
page = 1
|
||||
|
||||
while True:
|
||||
if args.max_pages is not None and page > args.max_pages:
|
||||
break
|
||||
|
||||
payload = fetch_page(fetch_args, page)
|
||||
scanned_pages += 1
|
||||
works = payload.get("works", [])
|
||||
if not isinstance(works, list) or not works:
|
||||
break
|
||||
|
||||
page_new = 0
|
||||
for item in works:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
keys = identity_keys(item)
|
||||
if not keys:
|
||||
continue
|
||||
if any(key in existing_keys for key in keys):
|
||||
for key in keys:
|
||||
fetched_existing_by_key[key] = item
|
||||
continue
|
||||
if any(key in seen_new_keys for key in keys):
|
||||
continue
|
||||
new_works.append(item)
|
||||
seen_new_keys.update(keys)
|
||||
page_new += 1
|
||||
|
||||
print(
|
||||
f"Fetched latest page {page}: {page_new} new / {len(works)} works",
|
||||
file=sys.stderr,
|
||||
)
|
||||
if page_new == 0:
|
||||
break
|
||||
page += 1
|
||||
|
||||
return new_works, fetched_existing_by_key, scanned_pages
|
||||
|
||||
|
||||
def merge_raw_works(
|
||||
existing_works: list[dict[str, Any]],
|
||||
new_works: list[dict[str, Any]],
|
||||
fetched_existing_by_key: dict[str, dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
merged: list[dict[str, Any]] = []
|
||||
emitted_keys: set[str] = set()
|
||||
|
||||
for work in new_works:
|
||||
keys = identity_keys(work)
|
||||
if any(key in emitted_keys for key in keys):
|
||||
continue
|
||||
merged.append(work)
|
||||
emitted_keys.update(keys)
|
||||
|
||||
for work in existing_works:
|
||||
replacement = next((fetched_existing_by_key[key] for key in identity_keys(work) if key in fetched_existing_by_key), None)
|
||||
candidate = replacement or work
|
||||
keys = identity_keys(candidate)
|
||||
if keys and any(key in emitted_keys for key in keys):
|
||||
continue
|
||||
merged.append(candidate)
|
||||
emitted_keys.update(keys)
|
||||
|
||||
return merged
|
||||
|
||||
|
||||
def numbered_works(works: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
result: list[dict[str, Any]] = []
|
||||
for index, work in enumerate(works, 1):
|
||||
numbered = dict(work)
|
||||
numbered["__line_no"] = index
|
||||
result.append(numbered)
|
||||
return result
|
||||
|
||||
|
||||
def filter_raw_works(args: argparse.Namespace, works: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], dict[str, int]]:
|
||||
exclude_tags = set(args.exclude_tag)
|
||||
exclude_tag_ids = set(args.exclude_tag_id)
|
||||
if args.keep_default_exclude_tags:
|
||||
exclude_tags.update(DEFAULT_EXCLUDE_TAGS)
|
||||
exclude_tag_ids.update(DEFAULT_EXCLUDE_TAG_IDS)
|
||||
return filter_works(numbered_works(works), exclude_tags, exclude_tag_ids)
|
||||
|
||||
|
||||
def read_json(path: Path) -> Any:
|
||||
with path.open(encoding="utf-8") as file:
|
||||
return json.load(file)
|
||||
|
||||
|
||||
def work_by_key(works: Iterable[dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
||||
result: dict[str, dict[str, Any]] = {}
|
||||
for work in works:
|
||||
for key in identity_keys(work):
|
||||
result.setdefault(key, work)
|
||||
return result
|
||||
|
||||
|
||||
def unique_by_primary_key(works: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
result: list[dict[str, Any]] = []
|
||||
seen: set[str] = set()
|
||||
for work in works:
|
||||
key = primary_key(work)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
result.append(work)
|
||||
return result
|
||||
|
||||
|
||||
def raw_diff_reasons(old: dict[str, Any], new: dict[str, Any]) -> list[str]:
|
||||
old_normalized = normalize_work(old, 0)
|
||||
new_normalized = normalize_work(new, 0)
|
||||
checks = (
|
||||
("title", old_normalized.get("title"), new_normalized.get("title")),
|
||||
("tags", old_normalized.get("tags"), new_normalized.get("tags")),
|
||||
("vas", old_normalized.get("vas"), new_normalized.get("vas")),
|
||||
("circle", old_normalized.get("circle"), new_normalized.get("circle")),
|
||||
("duration", old_normalized.get("duration"), new_normalized.get("duration")),
|
||||
("dlCount", old_normalized.get("dlCount"), new_normalized.get("dlCount")),
|
||||
("rateAverage", old_normalized.get("rateAverage"), new_normalized.get("rateAverage")),
|
||||
("rateCount", old_normalized.get("rateCount"), new_normalized.get("rateCount")),
|
||||
("release", old_normalized.get("release"), new_normalized.get("release")),
|
||||
("createDate", old_normalized.get("createDate"), new_normalized.get("createDate")),
|
||||
)
|
||||
return [name for name, old_value, new_value in checks if old_value != new_value]
|
||||
|
||||
|
||||
def updated_existing_diffs(
|
||||
existing_works: list[dict[str, Any]],
|
||||
fetched_existing_by_key: dict[str, dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
existing_by_key = work_by_key(existing_works)
|
||||
diffs: list[dict[str, Any]] = []
|
||||
for new_work in unique_by_primary_key(fetched_existing_by_key.values()):
|
||||
old_work = next((existing_by_key[key] for key in identity_keys(new_work) if key in existing_by_key), None)
|
||||
if old_work is None:
|
||||
continue
|
||||
reasons = raw_diff_reasons(old_work, new_work)
|
||||
if reasons:
|
||||
diffs.append({"work": new_work, "reasons": reasons})
|
||||
return diffs
|
||||
|
||||
|
||||
def filtered_change_diffs(
|
||||
filtered_path: Path,
|
||||
filtered_works: list[dict[str, Any]],
|
||||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
old_filtered = read_jsonl(filtered_path) if filtered_path.exists() else []
|
||||
old_by_key = work_by_key(old_filtered)
|
||||
new_by_key = work_by_key(filtered_works)
|
||||
|
||||
added: list[dict[str, Any]] = []
|
||||
removed: list[dict[str, Any]] = []
|
||||
emitted_added: set[str] = set()
|
||||
emitted_removed: set[str] = set()
|
||||
for work in filtered_works:
|
||||
key = primary_key(work)
|
||||
if key not in old_by_key and key not in emitted_added:
|
||||
added.append(work)
|
||||
emitted_added.add(key)
|
||||
for work in old_filtered:
|
||||
key = primary_key(work)
|
||||
if key not in new_by_key and key not in emitted_removed:
|
||||
removed.append(work)
|
||||
emitted_removed.add(key)
|
||||
return added, removed
|
||||
|
||||
|
||||
def embedding_diff_reasons(old: dict[str, Any], new: dict[str, Any]) -> list[str]:
|
||||
reasons: list[str] = []
|
||||
for field in ("title", "tags", "vas"):
|
||||
if old.get(field) != new.get(field):
|
||||
reasons.append(field)
|
||||
return reasons or ["embedding text"]
|
||||
|
||||
|
||||
def jsonl_record_count(path: Path) -> int:
|
||||
if not path.exists():
|
||||
return -1
|
||||
count = 0
|
||||
with path.open(encoding="utf-8") as file:
|
||||
for line in file:
|
||||
if line.strip():
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def web_data_out_of_sync(filtered_path: Path, data_dir: Path) -> bool:
|
||||
manifest_path = data_dir / "manifest.json"
|
||||
if not manifest_path.exists():
|
||||
return True
|
||||
try:
|
||||
manifest = read_json(manifest_path)
|
||||
web_count = int(manifest.get("count", -1))
|
||||
except (OSError, ValueError, TypeError, json.JSONDecodeError):
|
||||
return True
|
||||
return jsonl_record_count(filtered_path) != web_count
|
||||
|
||||
|
||||
def old_embedding_index_by_key(works: list[dict[str, Any]]) -> dict[str, int]:
|
||||
index_by_key: dict[str, int] = {}
|
||||
for index, work in enumerate(works):
|
||||
if not isinstance(work, dict):
|
||||
continue
|
||||
for key in identity_keys(work):
|
||||
index_by_key.setdefault(key, index)
|
||||
return index_by_key
|
||||
|
||||
|
||||
def existing_embedding_index(work: dict[str, Any], index_by_key: dict[str, int]) -> int | None:
|
||||
for key in identity_keys(work):
|
||||
index = index_by_key.get(key)
|
||||
if index is not None:
|
||||
return index
|
||||
return None
|
||||
|
||||
|
||||
def request_missing_embeddings(
|
||||
*,
|
||||
texts: list[str],
|
||||
start_indexes: list[int],
|
||||
manifest: dict[str, Any],
|
||||
args: argparse.Namespace,
|
||||
) -> dict[int, bytes]:
|
||||
if not texts:
|
||||
return {}
|
||||
|
||||
base_url = manifest.get("baseUrl")
|
||||
model = manifest.get("model")
|
||||
if not base_url or not model:
|
||||
raise RuntimeError("manifest must contain baseUrl and model for remote embeddings")
|
||||
|
||||
batch_size = int(manifest.get("batchSize") or 32)
|
||||
concurrency = int(manifest.get("concurrency") or 4)
|
||||
retry_forever = bool(manifest.get("retryForever", True))
|
||||
max_retry_wait = float(manifest.get("maxRetryWait") or 60.0)
|
||||
should_normalize = bool(manifest.get("normalized", True))
|
||||
url = embeddings_url(str(base_url))
|
||||
api_key = os.environ.get(args.api_key_env, "")
|
||||
packed_by_index: dict[int, bytes] = {}
|
||||
pending = {}
|
||||
completed: dict[int, list[list[float]]] = {}
|
||||
next_submit = 0
|
||||
next_write = 0
|
||||
dimensions = int(manifest["dimensions"])
|
||||
|
||||
def submit_available(executor: ThreadPoolExecutor) -> None:
|
||||
nonlocal next_submit
|
||||
while next_submit < len(texts) and len(pending) + len(completed) < concurrency:
|
||||
start = next_submit
|
||||
batch = texts[start : start + batch_size]
|
||||
future = executor.submit(
|
||||
request_embeddings_with_retries,
|
||||
url=url,
|
||||
model=str(model),
|
||||
inputs=batch,
|
||||
api_key=api_key,
|
||||
timeout=args.embedding_timeout,
|
||||
retries=args.embedding_retries,
|
||||
retry_wait=args.embedding_retry_wait,
|
||||
max_retry_wait=max_retry_wait,
|
||||
retry_forever=retry_forever,
|
||||
batch_start=start,
|
||||
)
|
||||
pending[future] = start
|
||||
next_submit += len(batch)
|
||||
|
||||
def pack_vectors(batch_start: int, vectors: list[list[float]]) -> None:
|
||||
for offset, vector in enumerate(vectors):
|
||||
if should_normalize:
|
||||
vector = normalize_vector(vector)
|
||||
if len(vector) != dimensions:
|
||||
raise RuntimeError(
|
||||
f"embedding dimensions mismatch: expected {dimensions}, got {len(vector)}"
|
||||
)
|
||||
packed_by_index[start_indexes[batch_start + offset]] = struct.pack(f"<{dimensions}f", *vector)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as executor:
|
||||
submit_available(executor)
|
||||
while pending:
|
||||
done, _ = wait(pending, return_when=FIRST_COMPLETED)
|
||||
for future in done:
|
||||
start = pending.pop(future)
|
||||
vectors = future.result()
|
||||
expected = min(batch_size, len(texts) - start)
|
||||
if len(vectors) != expected:
|
||||
raise RuntimeError(
|
||||
f"batch at {start} returned {len(vectors)} embeddings; expected {expected}"
|
||||
)
|
||||
completed[start] = vectors
|
||||
|
||||
while next_write in completed:
|
||||
vectors = completed.pop(next_write)
|
||||
pack_vectors(next_write, vectors)
|
||||
next_write += len(vectors)
|
||||
print(f"Embedded {next_write}/{len(texts)} missing works", file=sys.stderr, flush=True)
|
||||
|
||||
submit_available(executor)
|
||||
|
||||
return packed_by_index
|
||||
|
||||
|
||||
def write_embeddings_atomic(
|
||||
*,
|
||||
data_dir: Path,
|
||||
final_works: list[dict[str, Any]],
|
||||
old_works: list[dict[str, Any]],
|
||||
old_index_by_key: dict[str, int],
|
||||
old_embeddings_path: Path,
|
||||
missing_embeddings: dict[int, bytes],
|
||||
dimensions: int,
|
||||
) -> None:
|
||||
vector_size = dimensions * 4
|
||||
expected_size = len(old_works) * vector_size
|
||||
actual_size = old_embeddings_path.stat().st_size
|
||||
if actual_size != expected_size:
|
||||
raise RuntimeError(f"embedding size mismatch: expected {expected_size} bytes, got {actual_size} bytes")
|
||||
|
||||
temp_path = data_dir / "embeddings.f32.tmp"
|
||||
try:
|
||||
with old_embeddings_path.open("rb") as old_file, temp_path.open("wb") as new_file:
|
||||
with mmap.mmap(old_file.fileno(), 0, access=mmap.ACCESS_READ) as old_map:
|
||||
for index, work in enumerate(final_works):
|
||||
if index in missing_embeddings:
|
||||
new_file.write(missing_embeddings[index])
|
||||
continue
|
||||
old_index = existing_embedding_index(work, old_index_by_key)
|
||||
if old_index is None:
|
||||
raise RuntimeError(f"missing embedding for {primary_key(work)}")
|
||||
offset = old_index * vector_size
|
||||
new_file.write(old_map[offset : offset + vector_size])
|
||||
temp_path.replace(data_dir / "embeddings.f32")
|
||||
except Exception:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def build_incremental_web_data(args: argparse.Namespace, filtered_works: list[dict[str, Any]]) -> dict[str, int]:
|
||||
data_dir = Path(args.data_dir)
|
||||
manifest_path = data_dir / "manifest.json"
|
||||
works_path = data_dir / "works.json"
|
||||
embeddings_path = data_dir / "embeddings.f32"
|
||||
manifest = read_json(manifest_path)
|
||||
if manifest.get("method") != "remote-openai-compatible":
|
||||
raise RuntimeError("incremental web/data update currently supports only remote-openai-compatible data")
|
||||
|
||||
old_works = read_json(works_path)
|
||||
if not isinstance(old_works, list):
|
||||
raise RuntimeError(f"{works_path} must contain a JSON array")
|
||||
|
||||
dimensions = int(manifest["dimensions"])
|
||||
old_index_by_key = old_embedding_index_by_key(old_works)
|
||||
final_works: list[dict[str, Any]] = []
|
||||
missing_texts: list[str] = []
|
||||
missing_indexes: list[int] = []
|
||||
reused = 0
|
||||
|
||||
for raw_work in filtered_works:
|
||||
work = normalize_work(raw_work, len(final_works))
|
||||
old_index = existing_embedding_index(work, old_index_by_key)
|
||||
if old_index is None or embedding_text(work) != embedding_text(old_works[old_index]):
|
||||
missing_indexes.append(len(final_works))
|
||||
missing_texts.append(embedding_text(work))
|
||||
else:
|
||||
reused += 1
|
||||
final_works.append(work)
|
||||
|
||||
missing_embeddings = request_missing_embeddings(
|
||||
texts=missing_texts,
|
||||
start_indexes=missing_indexes,
|
||||
manifest=manifest,
|
||||
args=args,
|
||||
)
|
||||
write_embeddings_atomic(
|
||||
data_dir=data_dir,
|
||||
final_works=final_works,
|
||||
old_works=old_works,
|
||||
old_index_by_key=old_index_by_key,
|
||||
old_embeddings_path=embeddings_path,
|
||||
missing_embeddings=missing_embeddings,
|
||||
dimensions=dimensions,
|
||||
)
|
||||
|
||||
new_manifest = dict(manifest)
|
||||
new_manifest.update(
|
||||
{
|
||||
"count": len(final_works),
|
||||
"dimensions": dimensions,
|
||||
"embeddingFile": "embeddings.f32",
|
||||
"worksFile": "works.json",
|
||||
"generatedAt": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
)
|
||||
write_json_atomic(works_path, final_works)
|
||||
write_json_atomic(manifest_path, new_manifest, indent=2)
|
||||
return {"reused": reused, "embedded": len(missing_indexes), "count": len(final_works)}
|
||||
|
||||
|
||||
def dry_run_web_stats(args: argparse.Namespace, filtered_works: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
data_dir = Path(args.data_dir)
|
||||
manifest = read_json(data_dir / "manifest.json")
|
||||
if manifest.get("method") != "remote-openai-compatible":
|
||||
raise RuntimeError("incremental web/data update currently supports only remote-openai-compatible data")
|
||||
old_works = read_json(data_dir / "works.json")
|
||||
old_index_by_key = old_embedding_index_by_key(old_works)
|
||||
reused = 0
|
||||
embedded = 0
|
||||
embedding_diffs: list[dict[str, Any]] = []
|
||||
for raw_work in filtered_works:
|
||||
work = normalize_work(raw_work, reused + embedded)
|
||||
old_index = existing_embedding_index(work, old_index_by_key)
|
||||
if old_index is None:
|
||||
embedded += 1
|
||||
embedding_diffs.append({"work": work, "reason": "new work"})
|
||||
elif embedding_text(work) != embedding_text(old_works[old_index]):
|
||||
embedded += 1
|
||||
reasons = ", ".join(embedding_diff_reasons(old_works[old_index], work))
|
||||
embedding_diffs.append({"work": work, "reason": f"changed {reasons}"})
|
||||
else:
|
||||
reused += 1
|
||||
return {"reused": reused, "embedded": embedded, "count": reused + embedded, "embeddingDiffs": embedding_diffs}
|
||||
|
||||
|
||||
def limited_rows(rows: list[Any], limit: int) -> list[Any]:
|
||||
return rows if limit == 0 else rows[:limit]
|
||||
|
||||
|
||||
def remaining_count(rows: list[Any], limit: int) -> int:
|
||||
return 0 if limit == 0 else max(0, len(rows) - limit)
|
||||
|
||||
|
||||
def print_work_section(title: str, prefix: str, works: list[dict[str, Any]], limit: int) -> None:
|
||||
print(f"\n{title}: {len(works)}")
|
||||
for work in limited_rows(works, limit):
|
||||
print(f" {prefix} {format_work(work)}")
|
||||
remaining = remaining_count(works, limit)
|
||||
if remaining:
|
||||
print(f" ... {remaining} more")
|
||||
|
||||
|
||||
def print_updated_section(diffs: list[dict[str, Any]], limit: int) -> None:
|
||||
print(f"\nUpdated existing raw works: {len(diffs)}")
|
||||
for diff in limited_rows(diffs, limit):
|
||||
print(f" ~ {format_work(diff['work'])} / {', '.join(diff['reasons'])}")
|
||||
remaining = remaining_count(diffs, limit)
|
||||
if remaining:
|
||||
print(f" ... {remaining} more")
|
||||
|
||||
|
||||
def print_embedding_section(diffs: list[dict[str, Any]], limit: int) -> None:
|
||||
print(f"\nEmbedding requests: {len(diffs)}")
|
||||
for diff in limited_rows(diffs, limit):
|
||||
marker = "+" if diff["reason"] == "new work" else "~"
|
||||
print(f" {marker} {format_work(diff['work'])} / {diff['reason']}")
|
||||
remaining = remaining_count(diffs, limit)
|
||||
if remaining:
|
||||
print(f" ... {remaining} more")
|
||||
|
||||
|
||||
def print_dry_run_diff(
|
||||
*,
|
||||
args: argparse.Namespace,
|
||||
new_works: list[dict[str, Any]],
|
||||
existing_works: list[dict[str, Any]],
|
||||
fetched_existing_by_key: dict[str, dict[str, Any]],
|
||||
filtered_path: Path,
|
||||
filtered_works: list[dict[str, Any]],
|
||||
web_stats: dict[str, Any],
|
||||
) -> None:
|
||||
limit = args.diff_limit
|
||||
updated_diffs = updated_existing_diffs(existing_works, fetched_existing_by_key)
|
||||
filtered_added, filtered_removed = filtered_change_diffs(filtered_path, filtered_works)
|
||||
|
||||
print_work_section("New raw works", "+", new_works, limit)
|
||||
print_updated_section(updated_diffs, limit)
|
||||
print_work_section("Filtered additions", "+", filtered_added, limit)
|
||||
print_work_section("Filtered removals", "-", filtered_removed, limit)
|
||||
print_embedding_section(web_stats.get("embeddingDiffs", []), limit)
|
||||
|
||||
|
||||
def validate_args(args: argparse.Namespace) -> None:
|
||||
if args.page_size <= 0:
|
||||
raise SystemExit("--page-size must be greater than 0")
|
||||
if args.timeout <= 0:
|
||||
raise SystemExit("--timeout must be greater than 0")
|
||||
if args.retries < 0:
|
||||
raise SystemExit("--retries must be greater than or equal to 0")
|
||||
if args.max_pages is not None and args.max_pages <= 0:
|
||||
raise SystemExit("--max-pages must be greater than 0")
|
||||
if args.diff_limit < 0:
|
||||
raise SystemExit("--diff-limit must be greater than or equal to 0")
|
||||
if args.embedding_timeout <= 0:
|
||||
raise SystemExit("--embedding-timeout must be greater than 0")
|
||||
if args.embedding_retries < 0:
|
||||
raise SystemExit("--embedding-retries must be greater than or equal to 0")
|
||||
if args.embedding_retry_wait < 0:
|
||||
raise SystemExit("--embedding-retry-wait must be greater than or equal to 0")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
validate_args(args)
|
||||
input_path = Path(args.input)
|
||||
filtered_path = Path(args.filtered_output)
|
||||
if not input_path.exists():
|
||||
raise SystemExit(f"input file not found: {input_path}")
|
||||
|
||||
existing_works = read_jsonl(input_path)
|
||||
existing_keys = {key for work in existing_works for key in identity_keys(work)}
|
||||
new_works, fetched_existing_by_key, scanned_pages = fetch_latest_works(args, existing_keys)
|
||||
|
||||
if not new_works and not args.force and not web_data_out_of_sync(filtered_path, Path(args.data_dir)):
|
||||
print(f"Scanned pages: {scanned_pages}")
|
||||
print("No new works found. web/data update skipped.")
|
||||
return 0
|
||||
|
||||
merged_works = merge_raw_works(existing_works, new_works, fetched_existing_by_key)
|
||||
filtered_works, filter_stats = filter_raw_works(args, merged_works)
|
||||
web_stats = dry_run_web_stats(args, filtered_works) if args.dry_run else None
|
||||
|
||||
print(f"Scanned pages: {scanned_pages}")
|
||||
print(f"New raw works: {len(new_works)}")
|
||||
print(f"Raw works after merge: {len(merged_works)}")
|
||||
print(f"Filtered works: {filter_stats['output']} ({filter_stats['total_removed']} removed)")
|
||||
if web_stats:
|
||||
print(f"web/data would reuse {web_stats['reused']} embeddings and request {web_stats['embedded']} embeddings")
|
||||
print_dry_run_diff(
|
||||
args=args,
|
||||
new_works=new_works,
|
||||
existing_works=existing_works,
|
||||
fetched_existing_by_key=fetched_existing_by_key,
|
||||
filtered_path=filtered_path,
|
||||
filtered_works=filtered_works,
|
||||
web_stats=web_stats,
|
||||
)
|
||||
|
||||
if args.dry_run:
|
||||
print("Dry run: no files were written.")
|
||||
return 0
|
||||
|
||||
write_jsonl_atomic(input_path, merged_works)
|
||||
write_jsonl_atomic(filtered_path, filtered_works)
|
||||
web_stats = build_incremental_web_data(args, filtered_works)
|
||||
print(f"Updated {input_path}")
|
||||
print(f"Updated {filtered_path}")
|
||||
print(
|
||||
f"Updated {args.data_dir}: {web_stats['count']} works, "
|
||||
f"reused {web_stats['reused']} embeddings, requested {web_stats['embedded']} embeddings"
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user