update readme

Author: herilalaina

CITATION.cff

@@ -0,0 +1,69 @@
+cff-version: 1.2.0
+message: "If you use this software, please cite the paper below."
+title: "alphapfn: Fast Entropy Search via In-Context Learning"
+abstract: >-
+  Lightweight package for α-PFN, a Prior-Fitted Network that amortizes
+  information-theoretic acquisition functions (PES, MES, JES) for
+  Bayesian Optimization. α-PFN replaces the slow Monte-Carlo sampling
+  used by classical entropy-search methods with a single transformer
+  forward pass.
+type: software
+version: 0.0.1
+license: MIT
+repository-code: "https://github.com/automl/AlphaPFN"
+url: "https://github.com/automl/AlphaPFN"
+keywords:
+  - bayesian-optimization
+  - entropy-search
+  - prior-fitted-networks
+  - pfn
+  - in-context-learning
+  - amortized-inference
+  - acquisition-function
+authors:
+  - family-names: "Rakotoarison"
+    given-names: "Herilalaina"
+    email: "rakoheri@ad.helsinki.fi"
+    affiliation: "University of Helsinki"
+  - family-names: "Adriaensen"
+    given-names: "Steven"
+    email: "adriaens@cs.uni-freiburg.de"
+    affiliation: "University of Freiburg"
+  - family-names: "Viering"
+    given-names: "Tom"
+    email: "t.j.viering@tudelft.nl"
+    affiliation: "Delft University of Technology"
+  - family-names: "Hvarfner"
+    given-names: "Carl"
+    affiliation: "Meta"
+  - family-names: "Müller"
+    given-names: "Samuel"
+    affiliation: "Meta"
+  - family-names: "Hutter"
+    given-names: "Frank"
+    affiliation: "Prior Labs / ELLIS Institute Tübingen / University of Freiburg"
+  - family-names: "Bakshy"
+    given-names: "Eytan"
+    affiliation: "Meta"
+preferred-citation:
+  type: conference-paper
+  title: "α-PFN: Fast Entropy Search via In-Context Learning"
+  year: 2026
+  url: "https://openreview.net/forum?id=7Oonij8oLU"
+  conference:
+    name: "Forty-third International Conference on Machine Learning (ICML)"
+  authors:
+    - family-names: "Rakotoarison"
+      given-names: "Herilalaina"
+    - family-names: "Adriaensen"
+      given-names: "Steven"
+    - family-names: "Viering"
+      given-names: "Tom"
+    - family-names: "Hvarfner"
+      given-names: "Carl"
+    - family-names: "Müller"
+      given-names: "Samuel"
+    - family-names: "Hutter"
+      given-names: "Frank"
+    - family-names: "Bakshy"
+      given-names: "Eytan"

README.md

@@ -1,79 +1,103 @@
 # $\alpha$-PFN: Fast Entropy Search via In-Context Learning
 
-Lightweight package for $\alpha$-PFN, a Prior-Fitted Network for fast entropy search. Supported acquisitions: Predictive Entropy Search (PES), Max Value Entropy Search (MES), Joint-Entropy Search (JES).
+[![PyPI version](https://img.shields.io/pypi/v/alphapfn.svg)](https://pypi.org/project/alphapfn/)
+[![Python](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+[![Paper](https://img.shields.io/badge/Paper-OpenReview-b31b1b.svg)](https://openreview.net/forum?id=7Oonij8oLU)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/automl/AlphaPFN/blob/main/examples/quickstart.ipynb)
 
-> To reproduce our ICML paper experiments, see branch
+**$\alpha$-PFN** is a Prior-Fitted Network that amortizes information-theoretic acquisition functions for Bayesian Optimization. It replaces the slow Monte-Carlo sampling used by classical entropy-search methods with a single transformer forward pass — yielding **1.6× – 70× speed-ups** over fully-Bayesian GP entropy search while matching its optimization quality. Drop-in compatible with `botorch.optim.optimize_acqf`. Supported acquisitions: **PES**, **MES**, **JES**.
+
+<p align="center">
+  <img src="images/hero.gif" alt="Traditional GP-based Entropy Search samples optima via RFF and averages conditional entropies over N MC samples; α-PFN approximates the same acquisition in a single transformer forward pass.">
+</p>
+
+> To reproduce our ICML 2026 paper experiments, see branch
 > [`icml2026`](https://github.com/automl/AlphaPFN/tree/icml2026).
 
 ## Install
 
+```bash
+pip install "alphapfn[botorch]"
+```
+
+Or from source:
+
 ```bash
 git clone https://github.com/automl/AlphaPFN
 cd AlphaPFN
-uv sync
+uv sync --extra botorch
 ```
 
-PyPI release: TODO.
+Pretrained checkpoints (~20 MB) download automatically on the first `from_pretrained` call and cache under `~/.cache/alphapfn/`.
 
 ## Quick start
 
-A minimal BO loop driven by `botorch.optim.optimize_acqf`:
+A self-contained 2D BO loop on Branin, using `botorch.optim.optimize_acqf`:
 
 ```python
+import math
 import torch
-from botorch.optim import optimize_acqf  # needs botorch
+from botorch.optim import optimize_acqf
 from alphapfn import AlphaPFN
 
+# 1. Define the objective on the unit cube (α-PFN maximizes — we negate Branin).
+def branin(X):
+    x1 = 15.0 * X[..., 0] - 5.0
+    x2 = 15.0 * X[..., 1]
+    a, b, c = 1.0, 5.1 / (4 * math.pi**2), 5.0 / math.pi
+    r, s, t = 6.0, 10.0, 1.0 / (8 * math.pi)
+    return -(a * (x2 - b * x1**2 + c * x1 - r) ** 2
+             + s * (1 - t) * torch.cos(x1) + s)
+
+# 2. Initial design.
+torch.manual_seed(0)
+d, n_init, num_steps = 2, 5, 15
+X = torch.rand(n_init, d, dtype=torch.double)
+y = branin(X)
 bounds = torch.stack([torch.zeros(d), torch.ones(d)]).double()
+
+# 3. Load the pretrained acquisition; checkpoints download on first call.
 acqf = AlphaPFN.from_pretrained(acquisition="JES")
 
+# 4. BO loop.
 for step in range(num_steps):
-    # $\alpha$-PFN expects standardized targets.
-    y_std = (y - y.mean()) / (y.std() + 1e-8)
+    y_std = (y - y.mean()) / (y.std() + 1e-8)   # standardize targets
     acqf.fit(X, y_std)
-
-    X_next, acq_value = optimize_acqf(
-        acq_function=acqf,
-        bounds=bounds,
-        q=1,
-        num_restarts=5,
-        raw_samples=128,
-    )
-    y_next = objective(X_next.squeeze(0))
-    X = torch.cat([X, X_next.detach().double()], dim=0)
+    X_next, _ = optimize_acqf(acqf, bounds=bounds, q=1, num_restarts=5, raw_samples=128)
+    y_next = branin(X_next.squeeze(0))
+    X = torch.cat([X, X_next.detach().double()])
     y = torch.cat([y, y_next.detach().double().reshape(1)])
+    print(f"step {step+1:>2}: best so far = {y.max().item():.4f}")
 ```
 
-A runnable version lives at
-[`examples/bo_with_optimize_acqf.py`](examples/bo_with_optimize_acqf.py):
-
-```bash
-.venv/bin/python examples/bo_with_optimize_acqf.py --acquisition JES --steps 15
-```
+Runnable version: [`examples/bo_with_optimize_acqf.py`](examples/bo_with_optimize_acqf.py)
+or open the [Colab notebook](https://colab.research.google.com/github/automl/AlphaPFN/blob/main/examples/quickstart.ipynb).
 
 ## API
 
 ```python
 AlphaPFN.from_pretrained(
-    acquisition: str | None = None,   # or "MES", "JES"
+    acquisition: str | None = None,   # "PES" (default), "MES", or "JES"
     version: str = "v1",
     *,
     load_base_model: bool = False,
     ucb_beta: float = 2.0,
-    strict: bool = True,              # Pass strict=False to skip the input-range and standardization checks.
+    strict: bool = True,              # pass strict=False to skip input checks
 )
 ```
 
 Before fitting, prepare your data so that:
-  - You are maximizing. To minimize instead, negate your objective.
-    This is NOT checked, so forgetting it silently gives wrong results.
-  - Each input feature is rescaled to lie between 0 and 1.
-  - Targets are roughly standardized (subtract the mean, divide by the std).
+- **You are maximizing.** To minimize instead, negate your objective.
+  This is NOT checked, so forgetting it silently gives wrong results.
+- **Each input feature lies in `[0, 1]`.** Rescale your search space accordingly.
+- **Targets are roughly standardized** (subtract the mean, divide by the std).
 
+`strict=True` (default) validates the cube and standardization on every `fit`/`forward`; pass `strict=False` to skip.
 
-# Cite
+## Cite
 
-```latex
+```bibtex
 @inproceedings{
   rakotoarison2026alphapfn,
   title={{$\alpha$}-PFN: Fast Entropy Search via In-Context Learning},
@@ -82,4 +106,4 @@ Before fitting, prepare your data so that:
   year={2026},
   url={https://openreview.net/forum?id=7Oonij8oLU}
 }
-```
+```

alphapfn/__init__.py

@@ -8,4 +8,5 @@
 """
 from alphapfn.api import AlphaPFN, ALLOWED_ACQUISITIONS
 
-__all__ = ["AlphaPFN", "ALLOWED_ACQUISITIONS"]
+__version__ = "0.0.1"
+__all__ = ["AlphaPFN", "ALLOWED_ACQUISITIONS", "__version__"]

images/hero.gif

@@ -5,12 +5,34 @@ build-backend = "hatchling.build"
 [project]
 name = "alphapfn"
 version = "0.0.1"
-authors = [
-    {name = "ML Freiburg"},
-]
-description = "Fast Entropy Search via In-Context Learning"
+description = "Fast Entropy Search via In-Context Learning (PFN-based amortized acquisition functions for Bayesian Optimization)"
 readme = "README.md"
+license = { file = "LICENSE" }
 requires-python = ">=3.10"
+authors = [
+    { name = "Herilalaina Rakotoarison" },
+    { name = "Steven Adriaensen" },
+    { name = "Tom Viering" },
+    { name = "Carl Hvarfner" },
+    { name = "Samuel Müller" },
+    { name = "Frank Hutter" },
+    { name = "Eytan Bakshy" },
+]
+maintainers = [
+    { name = "Herilalaina Rakotoarison", email = "rakoheri@ad.helsinki.fi" },
+    { name = "Steven Adriaensen", email = "adriaens@cs.uni-freiburg.de" },
+    { name = "Tom Viering", email = "t.j.viering@tudelft.nl" },
+]
+keywords = [
+    "bayesian-optimization",
+    "entropy-search",
+    "prior-fitted-networks",
+    "pfn",
+    "in-context-learning",
+    "amortized-inference",
+    "acquisition-function",
+    "pytorch",
+]
 dependencies = [
     "torch>=1.9.0",
     "numpy>=1.21.2,<2.0",
@@ -23,17 +45,34 @@ dependencies = [
     "platformdirs",
 ]
 classifiers = [
-    "Programming Language :: Python :: 3",
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Mathematics",
 ]
 
+[project.urls]
+Homepage = "https://github.com/automl/AlphaPFN"
+Repository = "https://github.com/automl/AlphaPFN"
+Issues = "https://github.com/automl/AlphaPFN/issues"
+Paper = "https://openreview.net/forum?id=7Oonij8oLU"
+
 [project.optional-dependencies]
 botorch = ["botorch>=0.15.0"]
 
 [tool.hatch.build]
 exclude = [
     "*.gz",
     "*.pt",
+    "images/*",
+    "examples/*.ipynb",
 ]
 
 [tool.hatch.build.targets.wheel]