from __future__ import annotations
import argparse
import hashlib
import re
import sys
from pathlib import Path
import duckdb
ROW = "__arc_embed_row"
DEFAULT_VECTOR_COLUMN = "embedding"
MODEL_PARTS = ("tokenizer.json", "model.safetensors", "config.json")
MODEL_KEY_DOMAIN = b"subtoken/model-key/v1"
VERSION_RE = re.compile(
r"^subtoken (?P<build>\S+) \(model (?P<id>\S+)@(?P<revision>[0-9a-f]+), "
r"key (?P<key>[0-9a-f]+), dim (?P<dim>\d+)\)$"
)
class Refusal(Exception):
def sql_lit(s: str) -> str:
return "'" + s.replace("'", "''") + "'"
def sql_ident(s: str) -> str:
return '"' + s.replace('"', '""') + '"'
def model_key(model_id: str, revision: str, parts: dict[str, bytes]) -> str:
digest = hashlib.sha256()
digest.update(MODEL_KEY_DOMAIN)
digest.update(model_id.encode("utf-8"))
digest.update(b"\x00")
digest.update(revision.encode("utf-8"))
digest.update(b"\x00")
for name in MODEL_PARTS:
digest.update(parts[name])
return digest.hexdigest()
def parse_version(reported: str) -> dict[str, str]:
match = VERSION_RE.match(reported.strip())
if match is None:
raise Refusal(
f"the extension reported a version line this operator cannot read: "
f"{reported!r}. It expects `subtoken <build> (model <id>@<revision>, "
f"key <hex>, dim <n>)`; a line saying the model is unavailable means the "
f"artifact loaded but its weights did not."
)
return match.groupdict()
def split_release(release: str) -> tuple[str, str]:
model_id, sep, revision = release.rpartition("@")
if not sep or not model_id or not revision:
raise Refusal(
f"`model_release:` is {release!r}, which is not `<model-id>@<revision>`. "
f"It is the release the declared model directory was taken from, and the "
f"revision is the full commit the files came from — the extension hashes "
f"both into the address this operator checks against."
)
return model_id, revision
def check_declared_model(model_dir: Path, release: str, version: dict[str, str]) -> str:
declared_id, declared_revision = split_release(release)
parts: dict[str, bytes] = {}
for name in MODEL_PARTS:
path = model_dir / name
if not path.is_file():
raise Refusal(
f"the declared model asset cannot be checked against the extension: "
f"{path} is missing. The extension's model address covers "
f"{', '.join(MODEL_PARTS)}, so a directory carrying fewer than all "
f"three cannot reproduce it. Fetch the missing file from "
f"{declared_id} at {declared_revision}, or drop `model:` and "
f"`model_release:` and let the extension's own version line be the "
f"record of which model embedded this corpus."
)
parts[name] = path.read_bytes()
recomputed = model_key(declared_id, declared_revision, parts)
published = version["key"]
checked = recomputed[: len(published)]
if checked != published:
reported_release = f"{version['id']}@{version['revision']}"
raise Refusal(
f"the extension does not carry the model this Protocol declares. The "
f"Protocol declares {declared_id}@{declared_revision} in {model_dir}, "
f"which addresses to {checked}; the extension "
f"reports {reported_release} addressing to {published}. Either the "
f"declared directory holds different bytes from the ones compiled into "
f"the extension, or the extension derives its address a different way "
f"from the one this operator reproduces."
)
return recomputed
def main() -> int:
ap = argparse.ArgumentParser(description="Embed a text column into a vector column.")
ap.add_argument("--input", required=True, help="Parquet to read.")
ap.add_argument("--text-column", required=True, help="Column to embed.")
ap.add_argument("--extension", required=True, help="subtoken extension artifact.")
ap.add_argument("--out", required=True, help="Parquet to write.")
ap.add_argument(
"--vector-column",
default=DEFAULT_VECTOR_COLUMN,
help="Column to write the vector into.",
)
ap.add_argument("--model", help="Model directory to check the extension against.")
ap.add_argument("--model-release", help="`<model-id>@<revision>` that directory is.")
args = ap.parse_args()
src = Path(args.input)
if not src.is_file():
raise Refusal(f"the input Parquet {src} does not exist.")
extension = Path(args.extension)
if not extension.is_file():
raise Refusal(
f"the embedding extension {extension} is not on disk. It is a declared "
f"input, so a step in this Protocol has to put it there; this operator "
f"does not download one."
)
if bool(args.model) != bool(args.model_release):
raise Refusal(
"`model:` and `model_release:` are checked against each other and against "
"the extension, so neither means anything alone. Set both, or neither."
)
con = duckdb.connect(config={"allow_unsigned_extensions": "true"})
con.execute(f"LOAD {sql_lit(str(extension))}")
version = parse_version(con.execute("SELECT subtoken_version()").fetchone()[0])
dim = int(version["dim"])
if args.model:
check_declared_model(Path(args.model), args.model_release, version)
con.execute("SET threads TO 1")
con.execute("SET preserve_insertion_order TO true")
con.execute(
f"CREATE TABLE arc_src AS SELECT *, (row_number() OVER ()) - 1 AS {ROW} "
f"FROM read_parquet({sql_lit(str(src))})"
)
columns = [r[0] for r in con.execute("DESCRIBE arc_src").fetchall()]
vector_column = args.vector_column
clashes = [c for c in (vector_column,) if c in columns]
if columns.count(ROW) > 1:
clashes.append(ROW)
if clashes:
raise Refusal(
f"{src} already carries a column named {clashes[0]!r}; this operator adds "
f"{vector_column!r} and will not overwrite an input column. Name another "
f"with `vector_column:`."
)
if args.text_column not in columns:
raise Refusal(
f"{src} has no column {args.text_column!r}. It carries: "
f"{', '.join(c for c in columns if c != ROW)}."
)
rows = con.execute("SELECT count(*) FROM arc_src").fetchone()[0]
if not rows:
raise Refusal(f"{src} has no rows to embed.")
con.execute(
f"CREATE TABLE arc_vec AS SELECT {ROW}, "
f"subtoken_embed(coalesce(CAST({sql_ident(args.text_column)} AS VARCHAR), '')) AS vec "
f"FROM arc_src"
)
empty = con.execute(
"SELECT count(*) FROM arc_vec "
"WHERE list_sum(list_transform(vec, lambda x: abs(x))) = 0"
).fetchone()[0]
if empty:
print(
f"[text_embed] {empty} of {rows} rows tokenised to nothing in the "
f"model vocabulary and embed as zero vectors",
file=sys.stderr,
)
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
con.execute(
f"COPY (SELECT s.* EXCLUDE ({ROW}), "
f"CAST(v.vec AS FLOAT[{dim}]) AS {sql_ident(vector_column)} "
f"FROM arc_src s JOIN arc_vec v ON s.{ROW} = v.{ROW} ORDER BY s.{ROW}) "
f"TO {sql_lit(str(out))} (FORMAT parquet, COMPRESSION zstd)"
)
print(
f"[text_embed] {rows} rows · {dim}-d vectors → {vector_column} in {out} "
f"· {version['id']}@{version['revision']} key {version['key']}"
)
return 0
if __name__ == "__main__":
try:
sys.exit(main())
except Refusal as refusal:
print(f"text_embed: {refusal}", file=sys.stderr)
sys.exit(1)