Prompt Tuning

Demonstrates the four built-in PromptTuner implementations. All share the same hill-climbing loop:

1. Apply the current prompt to a labeled dataset
2. Evaluate accuracy (or another metric)
3. If not perfect, send the failing examples back to the model and ask for an improved prompt
4. Repeat until 100% accuracy (or the configured metric target) or max_iterations is reached

Section	Tuner	Task type
A-E	ClassificationPromptTuner	Binary YES/NO classification
F	MultiClassPromptTuner	N-way classification
G	ExtractionPromptTuner	Structured JSON field extraction
H	JudgedPromptTuner + LLMJudgeScorer	Open-ended generation rated by a judge model
I	JudgedPromptTuner + DeepEvalScorer	Open-ended generation rated by DeepEval metrics

Optionally persists each improvement to a PromptCatalog for version tracking.

A - Setup

Create a model client and build a small labeled dataset. The dataset must have: - content column: the text to classify - actual_class column: boolean ground truth

import pandas as pd
from aimu.prompts import ClassificationPromptTuner

# Use any BaseModelClient -- swap in AnthropicClient, OpenAIClient, etc.
# from aimu.models import AnthropicClient, AnthropicModel
# model_client = AnthropicClient(AnthropicModel.CLAUDE_SONNET_4_6)
from aimu.models import OllamaClient, OllamaModel

model_client = OllamaClient(OllamaModel.QWEN_3_5_9B)

# Labeled dataset: does this text discuss AI / machine learning?
df = pd.DataFrame(
    {
        "content": [
            "Large language models are trained on vast amounts of text data.",
            "The recipe calls for two cups of flour and one egg.",
            "Transformers revolutionised natural language processing in 2017.",
            "She planted tomatoes and basil in her garden this spring.",
            "Reinforcement learning from human feedback improves model alignment.",
            "The train departs from platform 3 at half past nine.",
            "Neural networks learn hierarchical representations from raw data.",
            "He spent the afternoon reading a novel by the lake.",
        ],
        "actual_class": [True, False, True, False, True, False, True, False],
    }
)

print(f"Dataset: {len(df)} rows, {df.actual_class.sum()} positive, {(~df.actual_class).sum()} negative")

B - Initial Prompt

Write a starting prompt. It must contain a {content} placeholder and instruct the model to respond with [YES] or [NO].

initial_prompt = (
    "Does the following text discuss artificial intelligence or machine learning?\n"
    "Answer with [YES] if it does, or [NO] if it does not.\n"
    "Only include [YES] or [NO] in your answer -- nothing else.\n"
    "Text: {content}"
)

tuner = ClassificationPromptTuner(model_client=model_client)

# Preview baseline performance before tuning
baseline_df = tuner.classify_data(initial_prompt, df)
baseline_metrics = tuner.evaluate_results(baseline_df)
print("Baseline metrics:")
for k, v in baseline_metrics.items():
    print(f"  {k}: {v:.3f}")

C - Run the Tuning Loop

tune() runs the hill-climbing loop. Each iteration that improves accuracy is logged to stdout. The loop stops when accuracy reaches 100% or max_iterations is exhausted.

best_prompt = tuner.tune(
    training_data=df,
    initial_prompt=initial_prompt,
    max_iterations=10,
    max_examples=4,
)

print("\n--- Best prompt found ---")
print(best_prompt)

D - Evaluate the Best Prompt

Run the best prompt on the full dataset and compare with the baseline.

final_df = tuner.classify_data(best_prompt, df)
final_metrics = tuner.evaluate_results(final_df)

print("Baseline vs. tuned:")
for k in ["accuracy", "precision", "recall"]:
    b = baseline_metrics[k]
    f = final_metrics[k]
    delta = f - b
    sign = "+" if delta >= 0 else ""
    print(f"  {k:12s}  baseline={b:.3f}  tuned={f:.3f}  ({sign}{delta:.3f})")

E - Catalog Persistence (Optional)

PromptCatalog stores each improved prompt version in a SQLite database so you can compare versions, roll back, and track metrics over time. Pass catalog and prompt_name to tune() to enable automatic saving.

import tempfile
import os
from aimu.prompts import PromptCatalog

db = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
db.close()

catalog = PromptCatalog(db.name)
prompt_name = "ai-classifier"

best_with_catalog = tuner.tune(
    training_data=df,
    initial_prompt=initial_prompt,
    max_iterations=10,
    catalog=catalog,
    prompt_name=prompt_name,
)

# Show all saved versions
model_id = str(model_client.model.value)
versions = catalog.retrieve_all(prompt_name, model_id)
print(f"Saved {len(versions)} prompt version(s) to catalog:")
for v in versions:
    acc = v.metrics.get("accuracy", "n/a") if v.metrics else "n/a"
    print(f"  v{v.version}  accuracy={acc}")
    print(f"  {v.prompt[:80]}..." if len(v.prompt) > 80 else f"  {v.prompt}")
    print()

os.unlink(db.name)

F - Multi-class Classification (MultiClassPromptTuner)

MultiClassPromptTuner extends the binary case to N named categories.

• Data must have a content column and an actual_class column (string, one of the registered classes).
• The prompt must instruct the model to output [ClassName] for one of the registered classes.
• Metrics include per-class precision, recall, F1, and macro F1.

import pandas as pd
from aimu.models import OllamaClient, OllamaModel
from aimu.prompts import MultiClassPromptTuner

model_client = OllamaClient(OllamaModel.QWEN_3_5_9B)

# Labeled dataset: 3-way sentiment classification
sentiment_df = pd.DataFrame(
    {
        "content": [
            "I absolutely loved this product, exceeded every expectation.",
            "Broken on arrival and customer support was useless.",
            "It does what it says. Nothing special, nothing bad.",
            "Best purchase I've made this year, highly recommend!",
            "Disappointing quality for the price. Would not buy again.",
            "Works fine. Exactly as described.",
            "Fantastic build quality and arrived ahead of schedule.",
            "Stopped working after two days. Very frustrating.",
            "Average product. Gets the job done but nothing more.",
        ],
        "actual_class": [
            "positive",
            "negative",
            "neutral",
            "positive",
            "negative",
            "neutral",
            "positive",
            "negative",
            "neutral",
        ],
    }
)

classes = ["positive", "negative", "neutral"]
multiclass_tuner = MultiClassPromptTuner(model_client=model_client, classes=classes)

# Preview baseline
initial_multiclass_prompt = (
    "Classify the sentiment of the following review.\n"
    "Reply with exactly one of: [positive], [negative], or [neutral].\n"
    "Review: {content}"
)

baseline_mc = multiclass_tuner.classify_data(initial_multiclass_prompt, sentiment_df)
baseline_mc_metrics = multiclass_tuner.evaluate_results(baseline_mc)

print("Baseline metrics:")
print(f"  accuracy:  {baseline_mc_metrics['accuracy']:.3f}")
print(f"  macro_f1:  {baseline_mc_metrics['macro_f1']:.3f}")
for cls in classes:
    f1 = baseline_mc_metrics[f"per_class_{cls}_f1"]
    print(f"  {cls:10s} f1={f1:.3f}")

best_multiclass_prompt = multiclass_tuner.tune(
    training_data=sentiment_df,
    initial_prompt=initial_multiclass_prompt,
    max_iterations=10,
    max_examples=4,
)

# Evaluate tuned prompt
final_mc = multiclass_tuner.classify_data(best_multiclass_prompt, sentiment_df)
final_mc_metrics = multiclass_tuner.evaluate_results(final_mc)

print("\nBaseline vs. tuned:")
for metric in ["accuracy", "macro_f1"]:
    b = baseline_mc_metrics[metric]
    f = final_mc_metrics[metric]
    sign = "+" if f >= b else ""
    print(f"  {metric:12s}  baseline={b:.3f}  tuned={f:.3f}  ({sign}{f - b:.3f})")

print("\n--- Best prompt ---")
print(best_multiclass_prompt)

G - Structured Field Extraction (ExtractionPromptTuner)

ExtractionPromptTuner tunes prompts that extract structured data from text.

• Data must have a content column and an expected column (dict mapping field name → expected value).
• The prompt must instruct the model to output JSON.
• A row is correct only if all specified fields match. Per-field accuracy is also tracked.
• The tuner handles raw JSON, fenced ```json ``` blocks, and any {...} substring automatically.

import pandas as pd
from aimu.models import OllamaClient, OllamaModel
from aimu.prompts import ExtractionPromptTuner

model_client = OllamaClient(OllamaModel.QWEN_3_5_9B)

# Labeled dataset: extract person name and employer from text
extraction_df = pd.DataFrame(
    {
        "content": [
            "Alice Smith is a senior engineer at Acme Corp.",
            "Bob Jones joined Initech as a product manager last year.",
            "Carol White leads the data science team at Globex.",
            "David Brown recently started his role as CTO at Umbrella Corp.",
            "Eva Green works as a software architect for Stark Industries.",
            "Frank Miller is employed by Wayne Enterprises as a security consultant.",
        ],
        "expected": [
            {"name": "Alice Smith", "company": "Acme Corp."},
            {"name": "Bob Jones", "company": "Initech"},
            {"name": "Carol White", "company": "Globex"},
            {"name": "David Brown", "company": "Umbrella Corp."},
            {"name": "Eva Green", "company": "Stark Industries"},
            {"name": "Frank Miller", "company": "Wayne Enterprises"},
        ],
    }
)

fields = ["name", "company"]
extraction_tuner = ExtractionPromptTuner(model_client=model_client, fields=fields)

initial_extraction_prompt = (
    "Extract the person's name and their employer from the text below.\n"
    'Return a JSON object with keys "name" and "company".\n'
    "Text: {content}"
)

# Preview baseline
baseline_ext_data = extraction_tuner.extract_data(initial_extraction_prompt, extraction_df)
baseline_ext_data = extraction_tuner.apply_prompt(initial_extraction_prompt, extraction_df)
baseline_ext_metrics = extraction_tuner.evaluate_results(baseline_ext_data)

print("Baseline metrics:")
print(f"  row accuracy:     {baseline_ext_metrics['accuracy']:.3f}")
for field in fields:
    print(f"  {field:10s} accuracy: {baseline_ext_metrics[f'field_{field}_accuracy']:.3f}")

best_extraction_prompt = extraction_tuner.tune(
    training_data=extraction_df,
    initial_prompt=initial_extraction_prompt,
    max_iterations=10,
    max_examples=3,
)

# Evaluate tuned prompt
final_ext_data = extraction_tuner.apply_prompt(best_extraction_prompt, extraction_df)
final_ext_metrics = extraction_tuner.evaluate_results(final_ext_data)

print("Baseline vs. tuned:")
for metric_key, label in [
    ("accuracy", "row accuracy"),
    ("field_name_accuracy", "name accuracy"),
    ("field_company_accuracy", "company accuracy"),
]:
    b = baseline_ext_metrics[metric_key]
    f = final_ext_metrics[metric_key]
    sign = "+" if f >= b else ""
    print(f"  {label:20s}  baseline={b:.3f}  tuned={f:.3f}  ({sign}{f - b:.3f})")

print("\n--- Best prompt ---")
print(best_extraction_prompt)

H - Open-ended Generation with a Scorer (JudgedPromptTuner)

JudgedPromptTuner handles tasks where there is no single correct answer (summarisation, rewriting, question answering) by delegating per-row evaluation to a pluggable Scorer. Two implementations ship with AIMU:

• LLMJudgeScorer (built-in): asks a separate judge model for a 1-10 rating and parses the score.
• DeepEvalScorer (aimu[evals]): wraps one or more DeepEval BaseMetric instances (e.g. GEval, AnswerRelevancyMetric, FaithfulnessMetric) and averages their scores.

Key differences from the other tuners: - Two model clients: model_client generates responses; the Scorer carries its own judge. - Primary metric is mean score (0-1), not binary accuracy. JudgedPromptTuner overrides score() to return metrics["score"] instead of metrics["accuracy"]. - pass_threshold: a row is _correct when judge_score >= pass_threshold (default 0.7). - Data uses a content column. An optional reference column is forwarded to scorers that consume it (used as LLMTestCase.expected_output for DeepEval metrics).

import pandas as pd
from aimu.models import OllamaClient, OllamaModel
from aimu.prompts import JudgedPromptTuner, LLMJudgeScorer

# Use the same or a different model as judge. A stronger model as judge is more reliable.
writer_client = OllamaClient(OllamaModel.QWEN_3_5_9B)
judge_client = OllamaClient(OllamaModel.QWEN_3_5_9B)

# Dataset: short articles to summarise.
# The reference column is optional; include it if you have a gold summary to compare against.
summarisation_df = pd.DataFrame(
    {
        "content": [
            (
                "Researchers at MIT have developed a new battery technology that can charge "
                "to 80% capacity in under five minutes. The breakthrough relies on a novel "
                "anode material made from silicon nanowires that dramatically increases ion "
                "transfer speed. Early prototypes have survived over 1,000 charge cycles "
                "without significant degradation, suggesting strong commercial potential for "
                "electric vehicles and grid storage."
            ),
            (
                "The city council voted 7-2 last Tuesday to approve a new zoning ordinance "
                "that will allow mixed-use development in the downtown core. Supporters argue "
                "the change will attract investment and reduce commuting by bringing housing "
                "closer to employment centres. Critics worry about gentrification and the loss "
                "of affordable units. The ordinance takes effect on the first of next month."
            ),
            (
                "A study published in Nature Medicine found that a Mediterranean diet "
                "supplemented with olive oil reduced the risk of cardiovascular events by 30% "
                "compared to a low-fat diet in a cohort of 7,400 participants followed over "
                "five years. The authors caution that dietary interventions are most effective "
                "when combined with regular physical activity and smoking cessation."
            ),
        ],
    }
)

criteria = (
    "A good summary is 1-2 sentences long, accurate, and covers the single most important "
    "finding or decision in the article. It should not include unnecessary detail."
)

# Built-in LLM-as-judge scorer: judge model rates each output 1-10.
scorer = LLMJudgeScorer(judge_client, criteria=criteria)

judged_tuner = JudgedPromptTuner(
    model_client=writer_client,
    scorer=scorer,
    pass_threshold=0.7,  # judge score >= 7/10 to be considered correct
)

initial_judged_prompt = "Summarise the following article in one or two sentences: {content}"

# Preview baseline scores
baseline_judged = judged_tuner.apply_prompt(initial_judged_prompt, summarisation_df)
baseline_judged_metrics = judged_tuner.evaluate_results(baseline_judged)

print("Baseline:")
print(f"  mean score:  {baseline_judged_metrics['score']:.3f}")
print(f"  pass rate:   {baseline_judged_metrics['pass_rate']:.3f}")
print()
for i, row in baseline_judged.iterrows():
    print(f"  [{row.judge_score:.1f}] {row.output[:100]}")

best_judged_prompt = judged_tuner.tune(
    training_data=summarisation_df,
    initial_prompt=initial_judged_prompt,
    max_iterations=8,
    max_examples=3,
)

# Evaluate tuned prompt
final_judged = judged_tuner.apply_prompt(best_judged_prompt, summarisation_df)
final_judged_metrics = judged_tuner.evaluate_results(final_judged)

print("Baseline vs. tuned:")
for label, key in [("mean score", "score"), ("pass rate", "pass_rate")]:
    b = baseline_judged_metrics[key]
    f = final_judged_metrics[key]
    sign = "+" if f >= b else ""
    print(f"  {label:12s}  baseline={b:.3f}  tuned={f:.3f}  ({sign}{f - b:.3f})")

print()
print("Tuned summaries:")
for i, row in final_judged.iterrows():
    print(f"  [{row.judge_score:.1f}] {row.output[:120]}")

print("\n--- Best prompt ---")
print(best_judged_prompt)

I - Open-ended Generation with DeepEval Metrics (DeepEvalScorer)

Swap LLMJudgeScorer for DeepEvalScorer to drive the same hill-climbing loop with one or more DeepEval metrics. Each metric is given an LLMTestCase built from the row (content -> input, generated output -> actual_output, optional reference -> expected_output); per-metric scores are averaged.

Requires aimu[evals].

from deepeval.metrics import GEval
from deepeval.test_case import LLMTestCaseParams

from aimu.evals import DeepEvalModel, DeepEvalScorer

# Reuse writer_client / judge_client / summarisation_df / criteria / initial_judged_prompt from section H.

geval = GEval(
    name="summary-quality",
    criteria=criteria,
    evaluation_params=[LLMTestCaseParams.INPUT, LLMTestCaseParams.ACTUAL_OUTPUT],
    model=DeepEvalModel(judge_client),
    threshold=0.7,
)

deepeval_tuner = JudgedPromptTuner(
    model_client=writer_client,
    scorer=DeepEvalScorer([geval]),
    pass_threshold=0.7,
)

best_deepeval_prompt = deepeval_tuner.tune(
    training_data=summarisation_df,
    initial_prompt=initial_judged_prompt,
    max_iterations=4,
    max_examples=3,
)

print("\n--- Best prompt (DeepEval-tuned) ---")
print(best_deepeval_prompt)