jytan12/causalean-retrieval

Causalean retrieval models

Two models that retrieve declarations from Causalean, a Lean 4 library formalizing causal inference (potential outcomes, structural causal models, identification, estimation theory, experimental design). They power the semantic tier of the library's search tool, which the CausalSmith research pipeline uses to find an existing lemma before proving a new one.

Path in this repo	Model	Base
/ (root)	bi-encoder retriever, 1024-dim, cosine similarity	BAAI/bge-large-en-v1.5
reranker/	cross-encoder reranker, one relevance score per (query, passage)	BAAI/bge-reranker-base

Both are task-specific fine-tunes: they are trained on, and meant for, this one library. They are not general-purpose mathematical retrievers.

Use inside the repository

You normally do not load these by hand. From a checkout of the repository:

scripts/fetch_retrieval_models.sh          # downloads this repository into doc/
cd CausalSmith/tools && npm run embed:library
npm run search -- --semantic "weak overlap rate for the ATE"

Set CAUSALEAN_MODELS_REV=<tag> to pin a revision (see Versioning). Without the models the tooling falls back to the off-the-shelf BAAI/bge-large-en-v1.5.

Use directly

Retriever — queries take the bge instruction prefix, passages do not:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("jytan12/causalean-retrieval")
prefix = "Represent this sentence for searching relevant passages: "
q = model.encode([prefix + "variance of the Horvitz-Thompson estimator under a Bernoulli design"],
                 normalize_embeddings=True)
p = model.encode(["Var htTotal. The design variance of the Horvitz-Thompson total equals ..."],
                 normalize_embeddings=True)
print(q @ p.T)

Reranker — it lives in a subfolder, so download it first:

from huggingface_hub import snapshot_download
from sentence_transformers import CrossEncoder

root = snapshot_download("jytan12/causalean-retrieval", allow_patterns="reranker/*")
reranker = CrossEncoder(f"{root}/reranker", max_length=256)
print(reranker.predict([["query text", "Declaration name. Its plain-English description."]]))

A passage is a declaration's humanized name followed by the first paragraph of its docstring (the plain-English translation of the formal statement), or its Lean statement when it has none.

Training data

Pairs mined from the library itself; nothing external.

• Retriever: contrastive pairs (MultipleNegativesRankingLoss) — a declaration's formal statement against its natural-language description, and a theorem's description against the descriptions of the declarations its proof uses, with hard negatives mined by the base model.
• Reranker: binary relevance (BinaryCrossEntropyLoss). A query is a theorem's description, a positive is a declaration its proof uses, and the negatives are the fine-tuned retriever's own top results for that query, which are the confusions a reranker has to resolve.

Model selection used held-out modules, so no declaration seen in training is scored at evaluation. Maximum sequence length is 256 tokens for both.

Versioning

The models describe one snapshot of the library and are retrained when the library's names or layout change materially. Each revision of this repository is tagged with the Causalean commit it was trained against; use the revision that matches your checkout, or the latest.

Limitations

English only. The passages are docstrings written for econometricians, so a query phrased as a Lean goal or in heavy notation retrieves worse than one phrased in words. Outside Causalean's vocabulary the models have no advantage over their base models.

License

Apache-2.0, the license of the repository. The base models are MIT-licensed.