certaintylabs 0.2.0

Python SDK for the Certainty Labs API — constraint enforcement for production LLMs

Certainty Python SDK

Python client for the Certainty Labs API — constraint enforcement for production LLMs.

Train TransEBM energy models and rerank LLM outputs in a few lines of code.

The SDK supports bring your own data (in-memory or local JSONL), tune training (epochs, batch size, model size, learning rate, etc.), and use your own LLM in rerank to generate candidates (openai_api_key + openai_base_url). You provide EORM-format training data; generate it externally if needed (see below).

Install

pip install certaintylabs

Quick Start

from certaintylabs import Certainty

client = Certainty()  # Uses fixed API URL; set CERTAINTY_API_KEY for auth

# Check the server is running
health = client.health()
print(health.version)  # "0.1.0"

# Train on the built-in GSM8K math reasoning dataset
result = client.train(epochs=10, d_model=768)
print(f"Accuracy: {result.best_val_acc:.1%} in {result.elapsed_seconds:.0f}s")

# Rerank LLM candidate answers
best = client.rerank(
    candidates=[
        "Janet sells 16 - 3 - 4 = 9 eggs. 9 * 2 = $18. The answer is 18.",
        "Janet has 16 eggs, sells all. 16 * 2 = $32.",
        "Janet sells 16 - 3 - 4 = 9 duck eggs. 9 * $2 = $18. The answer is $18.",
    ],
    prompt="Janet's ducks lay 16 eggs per day. She eats three and bakes muffins with four. She sells the rest at $2 each. How much does she make?",
)
print(best.best_candidate)

Async Support

from certaintylabs import AsyncCertainty

async def main():
    async with AsyncCertainty() as client:
        result = await client.train(epochs=5)
        best = await client.rerank(["A", "B", "C"], prompt="...")

Production: API key

The API base URL is fixed; you don't configure it. Set your API key via environment variable:

export CERTAINTY_API_KEY="ck_your_key_here"

from certaintylabs import Certainty

# Reads CERTAINTY_API_KEY from env
client = Certainty()
client.health()

You can also pass api_key explicitly to override the environment.

Logging

The SDK logs at INFO (client init, train, pipeline) and DEBUG (requests, responses). Enable verbose logs:

import logging
logging.getLogger("certaintylabs").setLevel(logging.DEBUG)

Data options

Option	SDK / API
Built-in dataset	train(epochs=10) with no data → uses GSM8K
Your data	train_with_data(samples) or train_from_file("path.jsonl") or train(data=...)
Rerank	rerank(candidates, prompt=...) or have the API generate candidates: rerank(prompt=..., openai_api_key=..., n_candidates=5)

Generating your own data externally

Training data must be EORM format: one JSON object per line with question, label (0 or 1), and gen_text. Create this data with your own pipeline (e.g. your LLM + your labeling rules or model-as-judge). Save as .jsonl and use train_from_file(path) or send the list to train(data=...).

API Reference

Certainty(api_key=None, timeout=300.0)

Parameter	Type	Default
api_key	str or None	None → env CERTAINTY_API_KEY or no auth
timeout	float	300.0

The API base URL is fixed and not configurable.

Methods

client.health() -> HealthResponse

Returns API status and version.

Using your own data

You can train on in-memory data or a local JSONL file instead of server-side data.

In-memory: each record is a dict with question, label, and gen_text (EORM format).

samples = [
    {"question": "What is 2+2?", "label": 1, "gen_text": "The answer is 4."},
    {"question": "What is 3*3?", "label": 1, "gen_text": "The answer is 9."},
]
result = client.train_with_data(samples, epochs=10)

Local file: one JSON object per line (same keys).

result = client.train_from_file("my_data.jsonl", epochs=15, lr=1e-4)

Low-level: pass data=... or data_path=... into client.train() for full control.

Tuning training parameters

Override defaults via keyword arguments or a TrainingParams object (omit fields to keep API defaults):

from certaintylabs import Certainty, TrainingParams

client = Certainty()

# Via kwargs
result = client.train(epochs=15, batch_size=2, lr=1e-4, max_length=1024)

# Via TrainingParams (good for reusing a config)
params = TrainingParams(epochs=15, batch_size=2, lr=1e-4, validate_every=2)
result = client.train(training_params=params)
# Or with your own data
result = client.train_with_data(samples, training_params=params)

TrainingParams supports: epochs, batch_size, d_model, n_heads, n_layers, lr, max_length, validate_every, val_holdout.

Rerank with your own model to generate candidates

You can either pass pre-generated candidates or have the API generate candidates with your LLM and then rerank them in one call. Use your own base model API (OpenAI, Claude, Llama, etc.) for generation:

# Option A: You provide candidates (e.g. from your own LLM elsewhere)
best = client.rerank(
    candidates=["answer A", "answer B", "answer C"],
    prompt="What is 2+2?",
)

# Option B: API generates n_candidates with your LLM, then reranks
best = client.rerank(
    prompt="What is 2+2?",
    openai_api_key="sk-...",
    openai_model="gpt-4o-mini",
    openai_base_url="https://api.openai.com/v1",
    n_candidates=5,
)
print(best.best_candidate)  # best of the 5 generated answers

client.train(**kwargs) -> TrainResponse

Train a TransEBM. Data source: data (list of records), data_path (server path), or neither (built-in GSM8K). Key parameters:

Parameter	Type	Default
yaml_content	str or None	None
data_path	str or None	None
data	list of dicts	None
epochs	int	20
d_model	int	768
n_heads	int	4
n_layers	int	2
lr	float	5e-5
max_length	int	2048
training_params	TrainingParams or None	None

result = client.train(data_path="path/to/gsm8k.jsonl", epochs=10)
print(result.model_path)     # "./certainty_workspace/model/..."
print(result.best_val_acc)   # 0.85

client.rerank(...) -> RerankResponse

Rerank LLM outputs using a trained TransEBM. Either pass candidates you already have, or omit candidates and set openai_api_key (and optionally openai_model, openai_base_url) so the API generates n_candidates with your LLM and then reranks them.

Parameter	Type	Default
candidates	List[str] or None	None (use with openai_api_key to generate)
prompt	str	""
model_path	str	"./certainty_workspace/model/ebm_certainty_model.pt"
tokenizer_path	str or None	None
openai_api_key	str or None	None
openai_model	str or None	None
openai_base_url	str or None	None
n_candidates	int	5 (only used when generating via your API)

best = client.rerank(
    candidates=["answer A", "answer B", "answer C"],
    prompt="What is 2+2?",
)
print(best.best_candidate)  # the highest-scored candidate
print(best.all_energies)    # energy scores for each candidate

client.pipeline(**kwargs) -> PipelineResponse

Run train (on your data or built-in) then optionally rerank. Pass data or data_path to use your data; omit for built-in. Pass candidates to rerank after training.

result = client.pipeline(epochs=10, candidates=["answer A", "answer B"])
print(result.train.best_val_acc)
if result.rerank:
    print(result.rerank.best_candidate)

Error Handling

from certaintylabs import Certainty, APIError, ConnectionError

client = Certainty()

try:
    client.compile("invalid yaml: [[[")
except APIError as e:
    print(e.status_code)  # 400
    print(e.detail)       # error message from the server

try:
    # ConnectionError when server unreachable (base URL is fixed)
    client = Certainty(timeout=2.0)
    client.health()
except ConnectionError as e:
    print(e)  # "Could not connect to <api-url>: ..."

License

MIT