mirror of
https://git.pumpin.online/funkpower/catena.git
synced 2026-10-22 18:57:14 +00:00
examples: programs that run, in four languages
Eight of them, each doing something a pipeline actually has to do, and each starting from something anyone can fetch — a public domain book from Project Gutenberg, a NASA broadcast, or a sample carried in the file. Shell, because the tools are filters first: a book to passages and what its notes cite; a recording to chapter marks, citable passages and the one ffmpeg call that makes it play on a phone; passages into pgvector and a search that asks twice and reconciles. Rust, for the case where the caller already holds the text and a subprocess would be serialising for nothing. Python, for the shape a pipeline stage takes once the reading is a library call. TypeScript, for the app side — chapter marks and deep links, and fusing two rankings in process. Three of them end with what re-ranking means, because it is two different things under one name: fusing ranked lists, which is what `rrf` does and needs no model, and a cross-encoder scoring query and passage together, which is accurate, slow, and runs over what retrieval already found.
This commit is contained in:
funkpower
@@ -0,0 +1,192 @@
|
||||
"""Passages into pgvector, then a search that asks twice and reconciles.
|
||||
|
||||
Loads the passages of a book into Postgres with both a vector column and a
|
||||
text index, searches both ways, and fuses the two rankings into one.
|
||||
|
||||
DATABASE_URL=postgresql://user@host/db \
|
||||
EMB_MODEL=nomic-embed-text \
|
||||
python examples/python/search_and_rerank.py "why did the empire decline"
|
||||
|
||||
Needs psycopg, a Postgres with pgvector, an embedding server with Ollama's
|
||||
API, and a book — the sample is fetched if none is given.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import subprocess
|
||||
import sys
|
||||
import urllib.request
|
||||
|
||||
import catena
|
||||
import psycopg
|
||||
|
||||
EMB_HOST = os.environ.get("EMB_HOST", "http://localhost:11434")
|
||||
EMB_MODEL = os.environ.get("EMB_MODEL", "nomic-embed-text")
|
||||
GUTENBERG = "https://www.gutenberg.org/cache/epub/731/pg731.epub"
|
||||
|
||||
|
||||
def embed(texts: list[str]) -> list[list[float]]:
|
||||
"""Ask the embedding server, in batches it will accept.
|
||||
|
||||
The one thing here that touches the network. No model is loaded in this
|
||||
process: it belongs behind a server, where it is shared and where its
|
||||
GPU is somebody else's problem.
|
||||
"""
|
||||
out: list[list[float]] = []
|
||||
for at in range(0, len(texts), 32):
|
||||
batch = texts[at : at + 32]
|
||||
body = json.dumps({"model": EMB_MODEL, "input": batch}).encode()
|
||||
request = urllib.request.Request(
|
||||
f"{EMB_HOST}/api/embed", data=body, headers={"content-type": "application/json"}
|
||||
)
|
||||
with urllib.request.urlopen(request, timeout=300) as reply: # noqa: S310
|
||||
out.extend(json.load(reply)["embeddings"])
|
||||
print(f" embedded {len(out)}/{len(texts)}", end="\r", file=sys.stderr)
|
||||
return out
|
||||
|
||||
|
||||
def passages(book: pathlib.Path) -> list[dict]:
|
||||
"""The book, cut into passages worth embedding."""
|
||||
markdown = subprocess.run(
|
||||
["pandoc", "-f", "epub", "-t", "commonmark", "--wrap=none", str(book)],
|
||||
capture_output=True, text=True, check=True,
|
||||
).stdout
|
||||
chunks = catena.chunk_markdown(catena.scrub(markdown))
|
||||
# A table of contents embeds as well as anything and retrieves for every
|
||||
# query, which is worse than not having it.
|
||||
return [c for c in chunks if c["role"] == "body" and c["tokens"] > 40]
|
||||
|
||||
|
||||
def load(conn, rows: list[dict], vectors: list[list[float]]) -> None:
|
||||
width = len(vectors[0])
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
cur.execute("DROP TABLE IF EXISTS catena_passages")
|
||||
# The tsvector is generated from the body rather than written
|
||||
# alongside it, so the two cannot drift apart.
|
||||
cur.execute(f"""
|
||||
CREATE TABLE catena_passages (
|
||||
idx integer PRIMARY KEY,
|
||||
chapter text,
|
||||
body text NOT NULL,
|
||||
fts tsvector GENERATED ALWAYS AS (to_tsvector('english', body)) STORED,
|
||||
vec vector({width})
|
||||
)""")
|
||||
with cur.copy("COPY catena_passages (idx, chapter, body, vec) FROM STDIN") as copy:
|
||||
for chunk, vector in zip(rows, vectors):
|
||||
copy.write_row((chunk["idx"], chunk["chapter"], chunk["text"], str(vector)))
|
||||
cur.execute("CREATE INDEX ON catena_passages USING gin (fts)")
|
||||
cur.execute("CREATE INDEX ON catena_passages USING hnsw (vec vector_cosine_ops)")
|
||||
conn.commit()
|
||||
|
||||
|
||||
def by_meaning(conn, vector: list[float], limit: int = 10) -> list[dict]:
|
||||
"""Nearest in the embedding space.
|
||||
|
||||
`<=>` is cosine distance, so smaller is nearer. This finds a passage
|
||||
about the fall of Rome for a query about decline, sharing no words.
|
||||
"""
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""SELECT idx, regexp_replace(left(body, 68), '\\s+', ' ', 'g')
|
||||
FROM catena_passages ORDER BY vec <=> %s::vector LIMIT %s""",
|
||||
(str(vector), limit),
|
||||
)
|
||||
return [
|
||||
{"v": 1, "id": str(idx), "channel": "vec", "rank": at, "score": 0.0, "text": body}
|
||||
for at, (idx, body) in enumerate(cur.fetchall(), start=1)
|
||||
]
|
||||
|
||||
|
||||
def by_word(conn, query: str, limit: int = 10) -> list[dict]:
|
||||
"""Nearest by the words themselves.
|
||||
|
||||
Finds the passage that actually says the word where the vector search
|
||||
drifted to a paraphrase, and finds nothing at all where the words are
|
||||
not in the book — which is a useful thing to be told.
|
||||
"""
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""SELECT idx, regexp_replace(left(body, 68), '\\s+', ' ', 'g')
|
||||
FROM catena_passages
|
||||
WHERE fts @@ websearch_to_tsquery('english', %(q)s)
|
||||
ORDER BY ts_rank_cd(fts, websearch_to_tsquery('english', %(q)s)) DESC
|
||||
LIMIT %(n)s""",
|
||||
{"q": query, "n": limit},
|
||||
)
|
||||
return [
|
||||
{"v": 1, "id": str(idx), "channel": "fts", "rank": at, "score": 0.0, "text": body}
|
||||
for at, (idx, body) in enumerate(cur.fetchall(), start=1)
|
||||
]
|
||||
|
||||
|
||||
def main() -> int:
|
||||
query = sys.argv[1] if len(sys.argv) > 1 else "why did the roman empire decline"
|
||||
url = os.environ.get("DATABASE_URL")
|
||||
if not url:
|
||||
print("set DATABASE_URL", file=sys.stderr)
|
||||
return 64
|
||||
|
||||
work = pathlib.Path("work")
|
||||
work.mkdir(exist_ok=True)
|
||||
book = work / "gibbon.epub"
|
||||
if not book.exists():
|
||||
print(f"fetching {GUTENBERG}")
|
||||
with urllib.request.urlopen(GUTENBERG, timeout=120) as reply: # noqa: S310
|
||||
book.write_bytes(reply.read())
|
||||
|
||||
rows = passages(book)
|
||||
print(f"== {len(rows)} passages")
|
||||
vectors = embed([" ".join(chunk["text"].split()) for chunk in rows])
|
||||
print(f" {EMB_MODEL} returns {len(vectors[0])} numbers per passage")
|
||||
|
||||
with psycopg.connect(url) as conn:
|
||||
load(conn, rows, vectors)
|
||||
print(" loaded into catena_passages")
|
||||
|
||||
print(f"\n== query: {query}")
|
||||
hits = by_meaning(conn, embed([query])[0]) + by_word(conn, query)
|
||||
|
||||
for channel, label in (("vec", "by meaning"), ("fts", "by word")):
|
||||
print(f"\n {label}:")
|
||||
for hit in [h for h in hits if h["channel"] == channel][:4]:
|
||||
print(f" {hit['rank']}. [{hit['id']}] {hit['text']}")
|
||||
|
||||
# Fusion is a library call rather than a subprocess: by the time a hit
|
||||
# reaches here it is already yours, and serialising it to hand to a
|
||||
# program and back is work for nothing.
|
||||
print("\n== fused")
|
||||
text_of = {hit["id"]: hit["text"] for hit in hits}
|
||||
for hit in catena.fuse_hits(hits, limit=5):
|
||||
print(f" {hit['rank']}. [{hit['id']}] score {hit['score']:.3f} {text_of[hit['id']]}")
|
||||
|
||||
print(RERANKING)
|
||||
return 0
|
||||
|
||||
|
||||
RERANKING = """
|
||||
Reciprocal rank fusion scores a passage by where it placed, not by what
|
||||
it scored: 1/(k + rank) from each list that found it, added up. The two
|
||||
searches return numbers that mean different things — a cosine distance
|
||||
and a ts_rank_cd — and no scaling makes them comparable, so their ranks
|
||||
are used and their scores thrown away. A passage both searches liked
|
||||
beats one that either loved, which is the behaviour worth having.
|
||||
|
||||
This is fusion, and it is often called re-ranking, but a second thing
|
||||
wears the same name. Retrieval uses a bi-encoder: passage and query are
|
||||
embedded separately, which is what lets the passage vectors be computed
|
||||
once and indexed, and is also what limits it — nothing in a passage's
|
||||
vector knows what was asked.
|
||||
|
||||
A cross-encoder reads the query and the passage together and scores the
|
||||
pair. It is far more accurate and far too slow to run over a corpus, so
|
||||
it runs over the twenty or fifty candidates retrieval already found. The
|
||||
two stack: retrieve both ways, fuse the lists, then hand the top of the
|
||||
fused list to a cross-encoder for the final order. Nothing here loads a
|
||||
model, so that step is yours — the fused list is what you would hand it.
|
||||
"""
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user