Skip to content

Getting Started

Requires Python >= 3.12.

pip install flash-ansr

This also pulls in symbolic-data and simplipy automatically, so no manual sequencing is needed. Check the installed version with flash_ansr.__version__.

Download a checkpoint

flash_ansr install psaegert/flash-ansr-v23.0-120M
By default models are cached under ./models/ relative to the package root and can be uninstalled with flash_ansr remove <repo>. Models can also be managed with the Python API via flash_ansr.model.manage.install_model and flash_ansr.model.manage.remove_model.

See all available models on Hugging Face:

Minimal inference Example

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Import flash_ansr
from flash_ansr import (
  FlashANSR,
  SoftmaxSamplingConfig,
  install_model,
  get_path,
)

# Select a model from Hugging Face
# https://huggingface.co/models?search=flash-ansr-v23.0
MODEL = "psaegert/flash-ansr-v23.0-120M"

# Download the latest snapshot of the model
# By default, the model is downloaded to the directory `./models/` in the package root
install_model(MODEL)

# Load the model (KV-cache, auto-batching and static decoding are on by default in v0.5)
model = FlashANSR.load(
  directory=get_path('models', MODEL),
  generation_config=SoftmaxSamplingConfig(choices=1024),  # or BeamSearchConfig / MCTSGenerationConfig
  length_penalty=0.05,  # prefer shorter expressions when scoring candidates (renamed from `parsimony` in v0.5)
).to(device)

# Define data
X = ...
y = ...

# Fit the model to the data
model.fit(X, y, verbose=True)

# Show the best expression
print(model.get_expression())

# Predict with the best expression
y_pred = model.predict(X)

Getting all candidates with infer

fit / get_expression / predict keep the fitted state on the model for read-back. To get every candidate in one call instead, use infer, which returns an InferenceResult and writes nothing to the model:

result = model.infer(X, y)

# Best refined candidate (or None if nothing fitted)
best = result.best
print(best.expression_infix)   # human-readable prediction
print(best.fvu, best.score, best.log_prob, best.constants)

# All refined survivors, score-sorted (best first)
for candidate in result.candidates:
    print(candidate.score, candidate.expression_infix)

# The full candidate ledger: the generation pool joined with the refined
# survivors, each classified FIT_OK / FIT_FAILED / INVALID
ledger = result.ledger
print(len(ledger))                       # total candidates considered
print(ledger.fit_status, ledger.fvu)     # per-candidate columns

# Timing of the two phases
print(result.generation_time, result.refinement_time)

# A tabular view of the refined survivors (one row per candidate in result.candidates)
df = result.to_dataframe()

A Candidate carries expression (skeleton tokens), expression_prefix, expression_infix, skeleton_prefix, constants, score, log_prob, fvu, complexity, constant_count, pruned_variant, and optional y_pred / y_pred_val (populated for the top top_k candidates). The FIT_OK / FIT_FAILED / INVALID codes live in flash_ansr.inference.

result.to_dataframe() returns a pandas DataFrame of the refined survivors (one row per candidate in result.candidates, i.e. FIT_OK fits), not the full ledger. To control which candidates get predictions, infer takes top_k (compute y_pred / y_pred_val for the top top_k candidates; None = the best only), predict_val (toggle validation-set prediction), and X_val (out-of-sample features for y_pred_val).

Find more details in the API Reference.

Evaluation

As of v0.6, evaluation, baseline comparisons, and benchmarking moved out of flash-ansr into the standalone srbf (Symbolic Regression Benchmark Framework) package.

pip install srbf

See the srbf repository for usage.

Next steps

  • See Concepts & Architecture for how the pieces fit together.
  • For training your own checkpoints, jump to Training.
  • For baseline comparisons and sweeps, see the srbf repository (evaluation moved to the standalone srbf package in v0.6).