125 lines
4.9 KiB
Python
125 lines
4.9 KiB
Python
"""Corpus embed + load entry point.
|
|
|
|
`cli.py` belongs to the parser/chunking work, so this is the standalone entry
|
|
point promised to Codex in `coordination/CLAUDE_TASK_2026-08-04.md` §4.1:
|
|
|
|
python -m ingestion.load.run --provider cohere-v4 --collection duocthu_v1
|
|
|
|
Two phases, deliberately separable. Embedding is the only part that leaves the
|
|
machine, so it goes through the disk cache: an interrupted run resumes from
|
|
whatever it already paid for instead of re-embedding it. Loading then reads
|
|
that cache and never calls a provider at all.
|
|
|
|
Vectors are keyed by `(model_id, input_kind, sha256(text))`, so a re-run after
|
|
an unrelated chunk edit re-embeds only the texts that actually changed.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Iterator, List, Sequence, Tuple
|
|
|
|
from ..embed.cache import CachingEmbeddingProvider, EmbeddingCache
|
|
from ..embed.ports import INPUT_DOCUMENT
|
|
from ..embed.registry import DEFAULT_REGION, build_provider
|
|
from .corpus import corpus_sha256, count_chunks, iter_chunk_records
|
|
from .models import CollectionSpec, CorpusManifest
|
|
from .qdrant_repo import QdrantVectorStore
|
|
from .upsert import ChunkLoader
|
|
|
|
DEFAULT_SLICE = 960
|
|
DEFAULT_ATTEMPTS = 3
|
|
|
|
|
|
def _slices(items: Sequence, size: int) -> Iterator[Tuple[int, Sequence]]:
|
|
for start in range(0, len(items), size):
|
|
yield start, items[start : start + size]
|
|
|
|
|
|
def _embed_with_retry(provider, texts: Sequence[str], attempts: int) -> List:
|
|
last: Exception | None = None
|
|
for attempt in range(1, attempts + 1):
|
|
try:
|
|
return list(provider.embed(texts, INPUT_DOCUMENT).vectors)
|
|
except Exception as exc: # noqa: BLE001 — a flaky link is the norm here
|
|
last = exc
|
|
if attempt == attempts:
|
|
break
|
|
backoff = 2**attempt
|
|
print(f" attempt {attempt} failed ({exc}); retrying in {backoff}s")
|
|
time.sleep(backoff)
|
|
raise RuntimeError(f"embedding failed after {attempts} attempts") from last
|
|
|
|
|
|
def main(argv=None) -> int:
|
|
parser = argparse.ArgumentParser(prog="python -m ingestion.load.run")
|
|
parser.add_argument("--chunks", type=Path, default=Path("data/processed/chunks.jsonl"))
|
|
parser.add_argument("--cache", type=Path, default=Path("data/processed/embeddings"))
|
|
parser.add_argument("--provider", required=True)
|
|
parser.add_argument("--collection", required=True)
|
|
parser.add_argument("--region", default=DEFAULT_REGION)
|
|
parser.add_argument("--qdrant-url", default="http://localhost:6333")
|
|
parser.add_argument("--slice-size", type=int, default=DEFAULT_SLICE)
|
|
parser.add_argument("--attempts", type=int, default=DEFAULT_ATTEMPTS)
|
|
parser.add_argument("--embed-only", action="store_true")
|
|
args = parser.parse_args(argv)
|
|
|
|
records = list(iter_chunk_records(args.chunks))
|
|
sha = corpus_sha256(args.chunks)
|
|
print(f"corpus : {args.chunks}")
|
|
print(f"chunks : {len(records)} (count_chunks={count_chunks(args.chunks)})")
|
|
print(f"sha256 : {sha}")
|
|
|
|
inner = build_provider(args.provider, region=args.region)
|
|
cache_path = args.cache / f"{args.provider}.jsonl"
|
|
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
|
provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
|
|
print(f"provider : {provider.model_id} ({provider.dimensions}d)")
|
|
print(f"cache : {cache_path}")
|
|
|
|
texts = [str(record["text"]) for record in records]
|
|
vectors: List = []
|
|
started = time.time()
|
|
for start, slice_texts in _slices(texts, args.slice_size):
|
|
vectors.extend(_embed_with_retry(provider, slice_texts, args.attempts))
|
|
done = len(vectors)
|
|
rate = done / max(time.time() - started, 1e-9)
|
|
remaining = (len(texts) - done) / rate if rate else 0
|
|
print(
|
|
f" embedded {done}/{len(texts)} "
|
|
f"({rate:.1f}/s, ~{remaining/60:.1f} min left)",
|
|
flush=True,
|
|
)
|
|
|
|
stats = provider.cache.stats
|
|
print(f"cache : {stats.hits} hits, {stats.misses} misses")
|
|
|
|
if args.embed_only:
|
|
print("embed-only: stopping before the vector store")
|
|
return 0
|
|
|
|
manifest = CorpusManifest(
|
|
corpus_sha256=sha,
|
|
chunk_count=len(records),
|
|
model_id=provider.model_id,
|
|
dimensions=provider.dimensions,
|
|
input_kind=INPUT_DOCUMENT,
|
|
provider=args.provider,
|
|
)
|
|
spec = CollectionSpec(name=args.collection, vector_size=provider.dimensions)
|
|
store = QdrantVectorStore(url=args.qdrant_url)
|
|
loader = ChunkLoader(store, spec, manifest)
|
|
|
|
report = loader.load(zip(records, (v.values for v in vectors)))
|
|
print(
|
|
f"loaded : {report.points_upserted} points in {report.batches} batches; "
|
|
f"collection holds {report.collection_count}; "
|
|
f"count gate {'PASS' if report.count_matches else 'FAIL'}"
|
|
)
|
|
return 0 if report.count_matches else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|