Skip to content

Repository files navigation

ANNEALER

Keep the prompt sharp as the science moves.

ANNEALER is a framework for continual prompt optimization of large language models in medicine. Where most prompt-optimization methods run once against a static benchmark, ANNEALER treats the system prompt as a living artifact: it periodically re-optimizes ("re-anneals") the prompt using textual gradients as new literature, guidelines, and benchmark items appear governed by a cooling schedule that controls how aggressively the prompt may change over time. The result is a generalist, open-weight model whose clinical reasoning tracks evolving evidence, with no weight updates and a full, auditable history of every prompt revision.

Key Features

  • Textual-gradient core - built on TextGrad-style automatic "differentiation" over text; optimizes the prompt, never the weights.
  • Cooling schedule - an annealing temperature trades exploration of new prompt edits against the stability of a proven prompt, so updates don't thrash.
  • Temporal evaluation - designed to be tested on time-stamped medical QA, measuring both how well it adapts to new evidence and how well it resists forgetting.
  • Evidence-aware - ingests fresh abstracts and guidelines as optimization signal, grounding each revision in current sources.
  • Training-free & open-weight - runs on models like Llama 3; no fine-tuning GPUs, fully reproducible.

The features above describe the design and aims of the framework. Benchmark numbers are populated after you run ANNEALER in your own environment.

How It Works

ANNEALER separates the model being improved from the optimizer that improves it:

  • Forward (student) model — the open-weight model under optimization (for example, Llama 3 70B) that answers each medical question using the current system prompt.
  • Backward (optimizer) engine — a separate engine that reads the student's mistakes, produces textual gradients (natural-language critiques), and proposes a revised system prompt.

Each optimization step runs the student over a batch of questions, scores the answers, backpropagates the textual gradients into the system-prompt variable, and applies an update. A held-out validation split gates every change: a revision that does not improve validation accuracy is reverted, and optimization stops early once improvements stall. The cooling schedule modulates how large a revision is accepted as optimization proceeds, and every accepted revision is logged so the full prompt history stays auditable.

Benchmarks

ANNEALER targets time-stamped medical QA, reporting both adaptation (accuracy on newly arrived evidence) and forgetting-resistance (retained accuracy on earlier items).

Benchmark Task Metric Result
MedQA USMLE-style multiple choice Accuracy TBD
PubMedQA Yes / No / Maybe over abstracts Accuracy TBD
NephSAP Nephrology multiple choice Accuracy TBD

Populate the table after running ANNEALER; the metrics above describe the evaluation plan rather than measured results.

Installation

# Clone your repository
git clone <your-repo-url>
cd annealer

# Install dependencies
pip install -r requirements.txt

Create a .env file in the project root with the API keys ANNEALER uses:

TOGETHER_API_KEY=...   # forward / student model (e.g. Llama 3 via Together)
OPENAI_API_KEY=...     # backward / optimizer engine

Usage

# Optimize a system prompt on a benchmark
python main.py --benchmark="MedQA4" --method="auto-prompt" --model="meta-llama/Llama-3-70b-chat-hf"

# Score a folder of result CSVs
python eval.py --csv_folder="path/to/results" --benchmark_name="NephSAP"

Supported --benchmark values: MedQA4, PubMedQA, NephSAP. Supported --method values: zero-shot, few-shot, CoT, auto-prompt.

Why Continual Matters

Most prompt-optimization methods optimize once on a static benchmark, so the prompt goes stale as evidence moves. ANNEALER's continual, temporal angle is built for fast-moving or sparsely-documented areas such as emerging therapeutics, shifting guidelines, and rare-disease subspecialties where evidence changes and labeled examples are scarce.

Project Layout

  • main.py — entry point; dispatches a benchmark/method/model run.
  • auto_prompt.py — wires data loaders to the prompt optimizer.
  • textgrad_wrapper.py — the optimization loop (forward pass, textual-gradient update, validation-gated revert, logging).
  • data.py, utils.py — dataset loading, batching, and answer evaluation.
  • mcq_baseline_methods.py, non_mcq_baseline_methods.py — zero-shot, few-shot, and chain-of-thought baselines.
  • run_post.py, eval.py, evaluation_scripts/ — post-hoc evaluation and scoring.
  • benchmarks/ — datasets (MedQA, PubMedQA, NephSAP).

License

Released under the MIT License. See LICENSE.

Acknowledgments

ANNEALER builds on TextGrad, a framework for optimizing LLM systems via textual gradients.

About

Continual, training-free prompt optimization for clinical LLMs which re-anneals system prompts with textual gradients as new medical evidence lands, so generalist models stay current without fine-tuning.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages