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,128 @@
|
||||
"""A book to rows a database can take, without a subprocess in sight.
|
||||
|
||||
This is the shape a pipeline stage ends up as once the reading is a library
|
||||
call: convert, then loop, and what is left is the part that is actually about
|
||||
your schema.
|
||||
|
||||
python examples/python/book_to_citations.py path/to/book.epub
|
||||
|
||||
With no argument it fetches Gibbon's first volume from Project Gutenberg,
|
||||
which is public domain and has a few hundred pages of footnotes.
|
||||
"""
|
||||
|
||||
import pathlib
|
||||
import subprocess
|
||||
import sys
|
||||
import urllib.request
|
||||
|
||||
import catena
|
||||
|
||||
GUTENBERG = "https://www.gutenberg.org/cache/epub/731/pg731.epub"
|
||||
|
||||
|
||||
def fetch(into: pathlib.Path) -> pathlib.Path:
|
||||
"""The sample book, downloaded once."""
|
||||
book = into / "gibbon.epub"
|
||||
if not book.exists():
|
||||
into.mkdir(parents=True, exist_ok=True)
|
||||
print(f"fetching {GUTENBERG}")
|
||||
with urllib.request.urlopen(GUTENBERG, timeout=120) as reply: # noqa: S310
|
||||
book.write_bytes(reply.read())
|
||||
return book
|
||||
|
||||
|
||||
def to_markdown(book: pathlib.Path) -> str:
|
||||
"""Pandoc is the one part of this that is somebody else's program.
|
||||
|
||||
`commonmark` rather than `commonmark_x`: the extended flavour writes
|
||||
anchors as bracketed spans, and they end up inside the passages.
|
||||
"""
|
||||
done = subprocess.run(
|
||||
["pandoc", "-f", "epub", "-t", "commonmark", "--wrap=none", str(book)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return done.stdout
|
||||
|
||||
|
||||
def rows(markdown: str) -> list[dict]:
|
||||
"""Passages in, citation rows out.
|
||||
|
||||
Two library calls do the work. `chunk_markdown` cuts the book into
|
||||
passages and flags the ones carrying citation material — which is the
|
||||
whole reason to ask it rather than looking for a heading, since a page of
|
||||
notes at the foot of a chapter has no heading over it. `find_citations`
|
||||
then reads each of those.
|
||||
"""
|
||||
out = []
|
||||
for chunk in catena.chunk_markdown(catena.scrub(markdown)):
|
||||
if "footnote_citations" not in chunk["flags"] and chunk["role"] != "bibliography":
|
||||
continue
|
||||
for found in catena.find_citations(chunk["text"]):
|
||||
# An `ibid.` or a short cite points at an entry above rather than
|
||||
# naming a work of its own, so there is nothing to store.
|
||||
if found["continuation"]:
|
||||
continue
|
||||
if not (found["title"] or found["authors"] or found["ids"]):
|
||||
continue
|
||||
out.append(
|
||||
{
|
||||
"page": chunk["page"],
|
||||
"chapter": chunk["chapter"],
|
||||
"citation_text": found["text"],
|
||||
"referenced_work": found["title"],
|
||||
"referenced_author": "; ".join(found["authors"]) or None,
|
||||
"referenced_year": found["year"],
|
||||
"referenced_container": found["container"],
|
||||
"kind": found["kind"],
|
||||
"unsettled": found["unsettled"],
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def main() -> int:
|
||||
where = pathlib.Path(sys.argv[1]) if len(sys.argv) > 1 else fetch(pathlib.Path("work"))
|
||||
|
||||
print(f"== {where.name}")
|
||||
identified = catena.identify(where.name, where.read_bytes())
|
||||
print(f" a {identified['kind']}, {identified['bytes']} bytes")
|
||||
|
||||
markdown = to_markdown(where)
|
||||
passages = catena.chunk_markdown(catena.scrub(markdown))
|
||||
carrying = [c for c in passages if "footnote_citations" in c["flags"]]
|
||||
print(f" {len(passages)} passages, {len(carrying)} carrying citation material")
|
||||
|
||||
found = rows(markdown)
|
||||
named = [r for r in found if r["referenced_work"]]
|
||||
print(f"\n== {len(found)} rows, {len(named)} naming a work")
|
||||
|
||||
# What a stage like this would do next is its own business: one executemany
|
||||
# against your own table. The point is that everything above it is two
|
||||
# function calls rather than a parser you have to keep.
|
||||
print("\n INSERT INTO citations (citation_text, referenced_work, referenced_author, …)")
|
||||
whole = [r for r in named if r["referenced_author"]] or named
|
||||
for row in whole[:5]:
|
||||
work = (row["referenced_work"] or "")[:40]
|
||||
who = (row["referenced_author"] or "—")[:22]
|
||||
print(f" (…, {work!r}, {who!r}, {row['referenced_year']})")
|
||||
|
||||
# Nothing is invented. Where a field could not be read the record says so
|
||||
# rather than guessing, and a caller that wants a model to fill the gaps
|
||||
# has `citation_prompt` and `merge_citation_reply` for exactly that.
|
||||
# Gibbon cites in the classical short form — `Tacit. Annal. xiii. 5.` —
|
||||
# so most of his notes name a passage rather than a book. Nothing is
|
||||
# invented to cover that: the record says which fields it could not read,
|
||||
# and `citation_prompt` and `merge_citation_reply` are there for a caller
|
||||
# that wants a model to finish them.
|
||||
unsettled = [r for r in found if r["unsettled"]]
|
||||
print(f"\n {len(unsettled)} rows name something they could not finish reading:")
|
||||
for row in unsettled[:3]:
|
||||
print(f" {row['citation_text'][:56]} → {', '.join(row['unsettled'])}")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -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