kuhyx
71e4a6115e
Merges ~/ck3-knowledge and ~/kcd-knowledge, which were 90% duplicate code and
neither of which had ever been committed anywhere.
Works on any wiki, not just MediaWiki:
- sources.py splits the engine-specific parts (enumerate, fetch) behind a Source
protocol. MediaWikiSource is preferred whenever an api.php exists;
SitemapHtmlSource handles custom engines with no API (verified on Fextralife).
Everything downstream was already source-agnostic.
- scripts/detect_wiki.py probes /api.php, /w/api.php then /sitemap.xml and prints
the config to use. The API path and article URL vary independently — Fandom
serves /api.php but /wiki/{Title}.
- Profiles are plain env files, since config.py was already fully env-driven.
Per-wiki behaviour that must be measured rather than assumed:
- Subpage policy. On the CK3 wiki "A/B" is a mod; on the KCD wiki it is the
sequel, and rejecting subpages deletes 441 pages of wanted content.
- Noise sections. "Other languages" tables were 7% of KCD chunks and matched
English queries by repeating every title in eight languages.
- Variant tagging for wikis covering two editions, with scopes that exclude the
other edition but keep both/unknown pages.
Build performance: length bucketing in embed_efficiently is 2.7x (fastembed pads
each batch to its longest member; 184-char median vs 2000-char max). With noise
sections that is 61 min -> 20 min on a 6,140-page wiki at half the workers.
Fandom portable-infoboxes are transformed to tables before markdownify —
otherwise stat values survive but detach from their labels.
100% branch coverage, ruff clean.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
61 lines
2.1 KiB
Python
61 lines
2.1 KiB
Python
"""Local embeddings via fastembed (bge-small), L2-normalized for cosine = dot."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable
|
|
from typing import Protocol
|
|
|
|
import numpy as np
|
|
|
|
from . import config
|
|
|
|
|
|
class EmbeddingModel(Protocol):
|
|
"""Minimal fastembed-compatible surface (also satisfied by test fakes)."""
|
|
|
|
def embed(self, texts: Iterable[str], **kwargs: object) -> Iterable[np.ndarray]: ...
|
|
def query_embed(self, texts: Iterable[str], **kwargs: object) -> Iterable[np.ndarray]: ...
|
|
|
|
|
|
_model: EmbeddingModel | None = None
|
|
|
|
|
|
def get_model() -> EmbeddingModel:
|
|
"""Lazily construct the (heavy) fastembed model once."""
|
|
global _model
|
|
if _model is None: # pragma: no cover - exercised only with the real model
|
|
from fastembed import TextEmbedding
|
|
|
|
# threads caps onnxruntime's intra-op parallelism (the main CPU driver);
|
|
# CPUExecutionProvider avoids grabbing any accelerator.
|
|
_model = TextEmbedding(
|
|
model_name=config.EMBED_MODEL,
|
|
threads=config.EMBED_THREADS,
|
|
providers=["CPUExecutionProvider"],
|
|
)
|
|
return _model
|
|
|
|
|
|
def _normalize(matrix: np.ndarray) -> np.ndarray:
|
|
if matrix.size == 0:
|
|
return matrix.reshape(0, config.EMBED_DIM).astype(np.float32)
|
|
norms = np.linalg.norm(matrix, axis=1, keepdims=True)
|
|
norms[norms == 0] = 1.0
|
|
return (matrix / norms).astype(np.float32)
|
|
|
|
|
|
def embed_passages(texts: list[str], *, model: EmbeddingModel | None = None) -> np.ndarray:
|
|
"""Embed corpus passages -> (N, dim) normalized float32."""
|
|
model = model or get_model()
|
|
vectors = list(
|
|
model.embed(texts, batch_size=config.EMBED_BATCH_SIZE, parallel=config.EMBED_PARALLEL)
|
|
)
|
|
return _normalize(np.array(vectors, dtype=np.float32))
|
|
|
|
|
|
def embed_query(text: str, *, model: EmbeddingModel | None = None) -> np.ndarray:
|
|
"""Embed a search query -> (dim,) normalized float32 (with bge query prefix)."""
|
|
model = model or get_model()
|
|
vector = np.array(next(iter(model.query_embed([text]))), dtype=np.float32)
|
|
return _normalize(vector.reshape(1, -1))[0]
|