metafit 0.1.1__py3-none-any.whl

metafit/__init__.py

@@ -0,0 +1 @@
+"""MetaFit rework: agents, decision trees, and training."""

metafit/agents.py

@@ -0,0 +1,162 @@
+"""Text generation with local Hugging Face models or OpenAI-compatible APIs.
+
+Dependencies are imported only when their respective agent is initialized.
+Each call to ``generate`` is independent and returns the generated text.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class Agent(ABC):
+    """Shared interface for language-model agents."""
+
+    @abstractmethod
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: str | None = None,
+        **generation_kwargs: Any,
+    ) -> str:
+        """Return generated text for a prompt."""
+        pass
+
+
+class LocalAgent(Agent):
+    """Load a causal model with separate loading and generation options.
+
+    ``model_kwargs`` accepts dtype, device_map, quantization_config, revision,
+    and other model-loading options. ``tokenizer_kwargs`` configures tokenizer
+    loading. ``generation_kwargs`` and ``chat_template_kwargs`` store defaults
+    that can be overridden in ``generate``.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        temperature: float = 0.7,
+        *,
+        device: str | None = None,
+        max_new_tokens: int = 512,
+        model_kwargs: dict[str, Any] | None = None,
+        tokenizer_kwargs: dict[str, Any] | None = None,
+        generation_kwargs: dict[str, Any] | None = None,
+        chat_template_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        self.model_name = model_name
+        self.temperature = temperature
+        self.max_new_tokens = max_new_tokens
+        self.generation_kwargs = dict(generation_kwargs or {})
+        self.chat_template_kwargs = dict(chat_template_kwargs or {})
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, **(tokenizer_kwargs or {}))
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, **(model_kwargs or {}))
+        if (model_kwargs or {}).get("device_map") is None:
+            self.model.to(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        self.model.eval()
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: str | None = None,
+        *,
+        chat_template_kwargs: dict[str, Any] | None = None,
+        **generation_kwargs: Any,
+    ) -> str:
+        """Generate text, overriding stored generation and template defaults.
+
+        A zero temperature selects greedy decoding. Models with a chat template
+        receive formatted messages; other models receive plain text.
+        """
+        import torch
+
+        if self.tokenizer.chat_template:
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.append({"role": "user", "content": prompt})
+            text = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                **{**self.chat_template_kwargs, **(chat_template_kwargs or {})},
+            )
+            inputs = self.tokenizer(
+                text, return_tensors="pt", add_special_tokens=False
+            )
+        else:
+            text = f"{system_prompt}\n\n{prompt}" if system_prompt else prompt
+            inputs = self.tokenizer(text, return_tensors="pt")
+        inputs = inputs.to(self.model.device)
+
+        params = {
+            "max_new_tokens": self.max_new_tokens,
+            "temperature": self.temperature,
+            **self.generation_kwargs,
+            **generation_kwargs,
+        }
+        temperature = params.pop("temperature")
+        params.setdefault("do_sample", temperature > 0)
+        if params["do_sample"]:
+            params["temperature"] = temperature
+        pad_token_id = self.tokenizer.pad_token_id
+        if pad_token_id is None:
+            pad_token_id = self.tokenizer.eos_token_id
+        if pad_token_id is not None:
+            params.setdefault("pad_token_id", pad_token_id)
+
+        with torch.inference_mode():
+            output = self.model.generate(**inputs, **params)
+        sequences = output.sequences if hasattr(output, "sequences") else output
+        new_tokens = sequences[0, inputs["input_ids"].shape[-1]:]
+        return self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+
+
+class APIAgent(Agent):
+    """Generate text with OpenAI Chat Completions or a compatible endpoint.
+
+    ``client_kwargs`` accepts timeout, max_retries, and other client options.
+    ``generation_kwargs`` stores endpoint-specific defaults such as top_p,
+    max_completion_tokens, response_format, reasoning_effort, and extra_body.
+    Use ``temperature=None`` to omit temperature from requests.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model_name: str = "gpt-4",
+        temperature: float | None = 0.7,
+        *,
+        base_url: str | None = None,
+        client_kwargs: dict[str, Any] | None = None,
+        generation_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        """Use the supplied key, or OPENAI_API_KEY when it is omitted."""
+        from openai import OpenAI
+
+        self.model_name = model_name
+        self.temperature = temperature
+        self.generation_kwargs = dict(generation_kwargs or {})
+        self.client = OpenAI(api_key=api_key, base_url=base_url, **(client_kwargs or {}))
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: str | None = None,
+        **generation_kwargs: Any,
+    ) -> str:
+        """Return one non-streaming reply, overriding stored generation defaults."""
+        messages = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+
+        params = {"temperature": self.temperature, **self.generation_kwargs, **generation_kwargs}
+        if params["temperature"] is None:
+            params.pop("temperature")
+        response = self.client.chat.completions.create(
+            model=self.model_name, messages=messages, **params
+        )
+        return response.choices[0].message.content

metafit/logger.py

@@ -0,0 +1,136 @@
+"""Training records, resume state, and analysis output in one place."""
+
+import hashlib
+import json
+from pathlib import Path
+
+from .tree import Candidate, LLMDecisionTree
+from .utils.analysis import (
+    save_config, save_json, save_training_log, save_final_result,
+    plot_training_progress, plot_test_performance,
+)
+
+
+class TrainerLogger:
+    def __init__(self, output_dir, resume=True):
+        self.output_dir = Path(output_dir)
+        self.resume = resume
+        self.history = []
+        self.candidates = []
+        self.cached_selection = None
+        self.agent_settings = {}
+        self.config = None
+
+    def save_inputs(self, config, split_indices):
+        """Record experiment configuration, seeds, and original row indices."""
+        self.agent_settings = {key: config[key] for key in [
+            "backend", "model_name", "temperature", "generation_kwargs",
+            "model_kwargs", "tokenizer_kwargs", "chat_template_kwargs",
+        ]}
+        self.config = config
+        self.split_indices = split_indices
+
+    def restore(self, trainer, batches):
+        """Check the saved footprint and rebuild completed trees, never refitting."""
+        self.history, self.candidates = [], []
+        self.cached_selection = None
+        settings = {key: getattr(trainer, key) for key in [
+            "random_state", "updates_per_epoch", "batch_size", "train_method",
+            "selection_metric", "regularizer_candidate_count", "stratify_batches",
+            "misclassified_top_k", "de_config", "data_config",
+        ]}
+        agent_settings = self.agent_settings or {
+            key: getattr(trainer.agent, key, None)
+            for key in ["model_name", "temperature", "generation_kwargs", "chat_template_kwargs"]
+        }
+        settings["agent_fingerprint"] = hashlib.sha256(
+            json.dumps(agent_settings, sort_keys=True, default=str).encode()
+        ).hexdigest()
+        for name, frame in [("train", trainer.train_df), ("val", trainer.val_df)]:
+            settings[name + "_fingerprint"] = hashlib.sha256(
+                frame.reset_index(drop=True).to_json(orient="table", double_precision=15).encode()
+            ).hexdigest()
+        settings = json.loads(json.dumps(settings))
+        path = self.output_dir / "training_state.json"
+        if self.resume and path.exists() and json.loads(path.read_text()) != settings:
+            raise ValueError("Training data or settings changed. Use a new output folder.")
+        save_json(path, settings)
+        if self.config is not None:
+            save_config(self.output_dir, self.config)
+            save_json(self.output_dir / "split.json", self.split_indices)
+        log_path = self.output_dir / "training_log.json"
+        if not self.resume or not log_path.exists():
+            return self.history, self.candidates
+
+        features = [c for c in trainer.train_df.columns if c != trainer.data_config["output_col"]]
+        for entry in json.loads(log_path.read_text(encoding="utf-8")):
+            restored = []
+            for record in entry["candidates"]:
+                tree = LLMDecisionTree(record["fitted_model"], feature_names=features)
+                tree.proposal_json = record["proposal"]
+                candidate = Candidate(
+                    tree, record["update"], record["candidate_index"],
+                    {"train": trainer.train_df if trainer.batch_size == 0 else batches[record["update"] - 1],
+                     "val": trainer.val_df, "full_train": trainer.train_df},
+                    trainer.data_config["output_col"],
+                )
+                for key in ["bootstrap_mean", "bootstrap_std", "bootstrap_valid_trials"]:
+                    setattr(candidate, key, record[key])
+                restored.append(candidate)
+            entry["candidate"] = next(c for c in restored if c.candidate_index == entry["selected_candidate_index"])
+            self.history.append(entry)
+            self.candidates.extend(restored)
+        selection = self.output_dir / "final_model_selection.json"
+        if selection.exists():
+            self.cached_selection = json.loads(selection.read_text(encoding="utf-8"))
+        return self.history, self.candidates
+
+    def log_update(self, trainer, update, prompt, feedback):
+        self.history.append({
+            "epoch": (update - 1) // trainer.updates_per_epoch + 1,
+            "batch": (update - 1) % trainer.updates_per_epoch + 1,
+            "update": update, "prompt": prompt,
+            "response": trainer.generation_attempts_[-1]["response"],
+            "attempts": trainer.generation_attempts_,
+            "candidate": trainer.current_candidate_, "sample_feedback": feedback,
+        })
+        self.save_training()
+
+    def save_training(self):
+        save_training_log(self.output_dir, self.history, self.candidates)
+
+    def selection_response(self, prompt, total_updates):
+        """Return a reusable reply, or None when the LLM must select again."""
+        cached = self.cached_selection
+        if cached and len(self.history) == total_updates and cached.get("prompt") == prompt:
+            return cached["response"]
+        return None
+
+    def save_selection(self, trainer):
+        selected = trainer.selected_candidate_
+        record = {
+            "method": trainer.return_weights, "metric": trainer.selection_metric,
+            "bootstrap_iterations": trainer.bootstrap_selection_iterations,
+            "bootstrap_sample_size": (trainer.bootstrap_selection_sample_size or len(trainer.val_df))
+            if trainer.bootstrap_selection_iterations > 0 else None,
+            "selected_candidate_rank": self.candidates.index(selected) + 1,
+            "selected_update": selected.update, "selected_candidate_index": selected.candidate_index,
+            "selected_bootstrap_mean": selected.bootstrap_mean,
+            "selected_bootstrap_std": selected.bootstrap_std,
+            "candidates": [
+                {"rank": rank, "update": c.update, "candidate_index": c.candidate_index,
+                 "validation_score": c.metrics["val"][trainer.selection_metric],
+                 "bootstrap_mean": c.bootstrap_mean, "bootstrap_std": c.bootstrap_std,
+                 "bootstrap_valid_trials": c.bootstrap_valid_trials}
+                for rank, c in enumerate(self.candidates, start=1)
+            ],
+        }
+        if trainer.return_weights == "llm_polish":
+            record.update(prompt=trainer.selection_prompt_, response=trainer.selection_response_)
+        save_json(self.output_dir / "final_model_selection.json", record)
+        return record
+
+    def save_result(self, result, metric="f1"):
+        save_final_result(self.output_dir, result)
+        plot_training_progress(self.output_dir, metric)
+        plot_test_performance(self.output_dir)

metafit/prompts/template.py

@@ -0,0 +1,306 @@
+"""Prompt templates and the prompt builder used during a MetaFit run."""
+
+from typing import Any
+
+import pandas as pd
+from ..tree import Candidate
+
+from ..utils.process import (
+    format_candidate_history,
+    format_feedback,
+    serialize_examples,
+    wrap_multiple_candidates,
+)
+
+
+def task_instructions(
+    output_col: str,
+    feature_names: list[str],
+    *,
+    task: str = "classification",
+    train_method: str = "exact_tree",
+    selection_metric: str = "f1",
+    hide_context: bool = False,
+) -> str:
+    """Describe the decision-tree task and misclassified-sample feedback."""
+    ml_model = "DecisionTree"
+    predictor_columns = feature_names
+    if hide_context:
+        metadata_guidance = f"""
+        NO-METADATA MODE IS ACTIVE.
+        The task context has intentionally been hidden. Do not try to infer, guess, recover, or name which real-world dataset this batch may come from.
+        Do not use external dataset memory, benchmark familiarity, domain priors, medical knowledge, UCI/OpenML knowledge, or any assumption about the source of these rows.
+        Even if feature names, class labels, feature ranges, or value patterns look familiar, treat them only as anonymous columns from an unknown tabular prediction task.
+        Fit the {ml_model} only from the provided labeled batch, the previous fitted model if present, and the explicit feedback signal. Use empirical relationships in these rows and feedback, not metadata.
+        If you recognize a possible dataset identity, ignore that recognition and continue as if the dataset identity is unknown."""
+    else:
+        metadata_guidance = "Think very carefully about what the predictors mean. Use your internal knowledge to estimate plausible model parameters or parameter ranges."
+
+    if train_method == "interval_tree":
+        proposal_description = "decision-tree structures whose thresholds are lower/upper ranges"
+        optimizer_guidance = f"""A numeric optimizer will choose exact thresholds within your ranges using a separate validation split and will target the selected metric: {selection_metric}. You do not see those validation rows. Propose ranges that should generalize to that metric rather than trying to reconstruct unseen validation cases."""
+    elif train_method == "exact_tree":
+        proposal_description = "decision-tree structures with exact numeric thresholds"
+        optimizer_guidance = f"""No numeric optimizer will alter your thresholds. Each exact tree will be evaluated directly on a separate validation split using the selected metric: {selection_metric}. You do not see those validation rows. Propose exact thresholds that should generalize to that metric rather than trying to reconstruct unseen validation cases."""
+    else:
+        raise ValueError(
+            f"Unsupported train_method={train_method!r}. "
+            "Choose 'exact_tree' or 'interval_tree'."
+        )
+
+    tabllm_example_note = """
+        The current batch below acts as few-shot labeled examples for this update.
+        The examples are serialized as natural-language feature-value statements.
+        Use them to understand the feature meanings and class-label patterns, but do not directly classify rows.
+        Return only the required JSON model."""
+    return f"""You are going to output {proposal_description} for a {ml_model} that predicts {output_col} from {predictor_columns}. It is a {task} problem.
+        {metadata_guidance}
+        {optimizer_guidance}
+        When feedback is provided, it contains the previous fitted model, grouped full-training-set error types, a scheduled number of high-confidence misclassified samples to fix, and the remaining feedback budget as high-confidence correct samples to preserve. For DecisionTree feedback, samples may also include the decision path taken through the fitted tree.
+        Use the previous fitted model as your starting point and adjust it conservatively. Prefer local corrections around its fitted thresholds over a completely new model. Change only the problematic subtree or thresholds suggested by the error groups and decision paths.
+        Treat high-confidence correct samples as do-no-harm constraints. Do not change branches or threshold regions that explain those correct predictions unless the grouped errors strongly justify it.
+        Use grouped error patterns to make structural corrections instead of overreacting to one row. Preserve useful thresholds; change them only when the error groups, misclassified-sample confidence, and decision paths suggest a specific correction.
+        Avoid overfitting to one misclassified sample, and do not infer broad changes from low-confidence errors alone.
+        Your output will be directly parsed by machine learning code, therefore make sure you output it correctly.
+        The current batch I am going to provide acts as few-shot examples for this update.
+        {tabllm_example_note}
+        You should not construct the tree entirely from the samples I provide to you.
+        You should follow very strictly to the output format that I am going to provide to you. 
+        Do not output anything else other than my given output format. Do not wrap the output in markdown, start with a bracket not '.
+        IMPORTANT: Output valid JSON only. Use lowercase true/false for boolean values, not Python-style True/False. Use double quotes for all strings and keys."""
+
+
+def base_prompt(
+    role_description: str,
+    task_description: str,
+    few_shot_examples: str,
+    *,
+    context_description: str = "",
+    feedback: str = "",
+    output_schema: str = "",
+) -> str:
+    """Assemble an update prompt; feedback already includes its heading."""
+    parts = [f"{role_description}\n", f"{task_description}\n"]
+    if context_description:
+        parts.append(f"Here is the context of the task. {context_description}\n")
+    if few_shot_examples:
+        parts.append(f"Here are some data from the task:\n{few_shot_examples}\n")
+    if feedback:
+        parts.append(feedback)
+    if output_schema:
+        parts.append(f"Here is the output format:{output_schema}\n")
+    return "\n".join(parts)
+
+
+def feedback_prompt(tree: str | None, misclassified_samples: str = "") -> str:
+    """Build the previous-attempt feedback section, or omit it when empty."""
+    feedback = format_feedback(tree, misclassified_samples)
+    if not feedback:
+        return ""
+    return f"""Here is the feedback from your previous attempt.
+
+                                Think critically on how you can improve the model based on the feedback, {feedback}\n"""
+
+
+def output_format(
+    task: str = "classification",
+    ml_model: str = "DecisionTree",
+    train_method: str = "exact_tree",
+) -> str:
+    """Return the decision-tree JSON output instructions for the strategy.
+
+    ``exact_tree`` requests exact parameters. ``interval_tree`` requests
+    parameter intervals. Template whitespace is preserved for compatibility.
+    """
+    formats = {
+        "classification": {
+            "DecisionTree": {
+                "format": """{
+                            "rationale": "<Brief model rationale: explain the feature-target relationships and causality used in the model. Use the error samples only to test this rationale; do not rely on them heavily.>",
+                            "tree": {
+                                "feature": "<feature_name>",
+                                "threshold": <threshold_value>,
+                                "left": {
+                                    "feature": "<feature_name>",
+                                    "threshold": <threshold_value>,
+                                    "left": {
+                                        "value": [<class_1>, <class_2>, ...]
+                                    },
+                                    "right": {
+                                        "value": [<class_1>, <class_2>, ...]
+                                    }
+                                },
+                                "right": {
+                                    "value": [<class_1>, <class_2>, ...]
+                                }
+                            },
+                            "classes": [<class_label_1>, <class_label_2>, ...]
+                        }
+
+                        *note: 
+                        - Internal (split) nodes MUST have "feature" and "threshold".
+                        - Leaf nodes MUST have "value": [<class_1>, <class_2>, ...] and NO "feature" or "threshold".
+                        - You can nest "left" and "right" as deeply as needed to fit the data.
+                        - For each leaf, set the predicted class to 1 and all others to 0.
+                        """,
+                "regularizer_format": """{
+                            "rationale": "<Brief model rationale: explain the feature-target relationships and causality used in the model. Use the error samples only to test this rationale; do not rely on them heavily.>",
+                            "tree": {
+                                "feature": "<feature_name>",
+                                "threshold": {"lower": <lower_bound>, "upper": <upper_bound>},
+                                "left": {
+                                    "feature": "<feature_name>",
+                                    "threshold": {"lower": <lower_bound>, "upper": <upper_bound>},
+                                    "left": {
+                                        "value": [<class_1>, <class_2>, ...]
+                                    },
+                                    "right": {
+                                        "value": [<class_1>, <class_2>, ...]
+                                    }
+                                },
+                                "right": {
+                                    "value": [<class_1>, <class_2>, ...]
+                                }
+                            },
+                            "classes": [<class_label_1>, <class_label_2>, ...]
+                        }
+
+                        *note: 
+                        - Internal (split) nodes MUST have "feature" and "threshold". The "threshold" must be an object with "lower" and "upper" bounds.
+                        - Leaf nodes MUST have "value": [<class_1>, <class_2>, ...] and NO "feature" or "threshold".
+                        - "classes" field must have names that exactly match the labels. Do not use descriptive name inside "classes" otherwise there would be an error.
+                        - You can nest "left" and "right" as deeply as needed to fit the data.
+                        - For each leaf, set the predicted class to 1 and all others to 0.
+                        - The "lower" and "upper" bounds define the acceptable range for the split threshold. The actual threshold will be aggressively optimized on the validation set. Make sure it will generalize well. If you see a perfect score, adjust the ranges to be narrower or prune the tree to prevent overfitting.
+                        """,
+            },
+        },
+    }
+    if train_method == "exact_tree":
+        format_name = "format"
+    elif train_method == "interval_tree":
+        format_name = "regularizer_format"
+    else:
+        raise ValueError(
+            f"Unsupported train_method={train_method!r}. "
+            "Choose 'exact_tree' or 'interval_tree'."
+        )
+    return formats[task][ml_model][format_name]
+
+def candidate_selection_prompt(
+    context: str,
+    history: str,
+    selection_metric: str = "f1",
+    feature_names: list[str] | None = None,
+) -> str:
+    """Ask the LLM to select one existing fitted candidate by displayed rank."""
+    feature_schema = ", ".join(feature_names) if feature_names else "Not available"
+    return f"""Select one final model from all exact fitted candidates below.
+    A Base ML reference candidate may also be included. Each candidate was evaluated on repeated bootstrap samples from the validation set; its bootstrap mean and standard deviation are reported for the selection metric ({selection_metric}), and full-training performance is included as additional fit context.
+    The candidate list can include both LLM-generated fitted models and a CV-tuned base ML model. Treat them fairly: do not assume an LLM candidate is better merely because it was generated by an LLM. Prefer candidates with high bootstrap mean and stable performance (low standard deviation), while using the dataset context only to resolve meaningful model-structure choices.
+
+    Dataset context:
+    {context if context else "Not available"}
+
+    Available feature names:
+    {feature_schema}
+
+    Selection procedure:
+    1. Treat train performance and validation performance as the primary evidence. High validation performance doesn't necessarily mean a good model, it could be overfitted on the validation set when train performance is significantly lower then validation performance. Reject models with big gap between train performance and val performance unless there is no any other choice.
+    2. Prefer a high bootstrap mean with a low bootstrap standard deviation.
+    3. Use full-training performance to identify unstable train-validation behavior.
+    4. When candidates have similar statistical evidence, prefer the simpler and more clinically plausible model.
+    5. Use dataset knowledge only as a tie-breaker. Do not ignore substantially better validation evidence merely because another model appears more plausible.
+
+    A perfect training or validation score may indicate overfitting, but it is not sufficient evidence by itself. Consider model complexity, train-validation agreement, bootstrap variability, and feature plausibility together.
+
+    Your only job is to choose exactly one existing candidate. Do not create, refine, merge, or modify a model. The selected candidate's stored fitted model will be returned unchanged. If models have similar train and validation performance, prefer models with training performance higher than the validation performance, and the gap between these performances should be small. Overall, your job is to select a model that generalizes the best.
+    Here are all exact fitted LLM candidates and the optional Base ML reference. Each listed tree is ready for inference and will be returned unchanged; no batch-training performance is shown: {history}.
+
+    Output only this JSON object, without markdown or explanation:
+    {{"selected_candidate_rank": <integer from 1 through the number of listed candidates>}}"""
+
+
+class PromptBuilder:
+    """Build DecisionTree prompts from fixed settings and per-call data.
+
+    Metadata requires ``output_col`` and ``role_description``. Optional keys
+    are ``task`` (default ``classification``), ``context_description``, and
+    ``hide_context``. Examples and feedback are never stored on the builder.
+    """
+
+    def __init__(
+        self,
+        metadata: dict[str, Any],
+        *,
+        train_method: str = "exact_tree",
+        selection_metric: str = "f1",
+        candidate_count: int = 5,
+        random_state: int = 114514,
+        rounding: int = 3,
+    ) -> None:
+        if candidate_count < 1:
+            raise ValueError("candidate_count must be positive.")
+        self.metadata = metadata.copy()
+        self.train_method = train_method
+        self.selection_metric = selection_metric
+        self.candidate_count = candidate_count
+        self.random_state = random_state
+        self.rounding = rounding
+
+    def build_update_prompt(
+        self,
+        examples: pd.DataFrame,
+        previous_tree: str | None = None,
+        sample_feedback: str = "",
+    ) -> str:
+        """Build a complete update prompt; an empty batch gives zero-shot input."""
+        output_col = self.metadata["output_col"]
+        task = self.metadata.get("task", "classification")
+        feature_names = [column for column in examples.columns if column != output_col]
+        examples_text = serialize_examples(
+            examples,
+            output_col,
+            random_state=self.random_state,
+            rounding=self.rounding,
+        )
+        instructions = task_instructions(
+            output_col,
+            feature_names,
+            task=task,
+            train_method=self.train_method,
+            selection_metric=self.selection_metric,
+            hide_context=bool(self.metadata.get("hide_context", False)),
+        )
+        prompt = base_prompt(
+            role_description=self.metadata["role_description"],
+            task_description=instructions,
+            few_shot_examples=examples_text,
+            context_description=self.metadata.get("context_description", ""),
+            feedback=feedback_prompt(previous_tree, sample_feedback),
+            output_schema=output_format(task=task, train_method=self.train_method),
+        )
+        return wrap_multiple_candidates(
+            prompt,
+            self.candidate_count,
+            deterministic=self.train_method == "exact_tree",
+        )
+
+    def build_selection_prompt(
+        self,
+        candidates: list[Candidate],
+        feature_names: list[str],
+    ) -> str:
+        """Describe candidates in their supplied display order and request a rank."""
+        history = format_candidate_history(candidates, self.selection_metric)
+        return candidate_selection_prompt(
+            context=self.metadata.get("context_description", ""),
+            history=history,
+            selection_metric=self.selection_metric,
+            feature_names=feature_names,
+        )
+
+    @staticmethod
+    def build_retry_prompt(original_prompt: str, retry_feedback: str) -> str:
+        """Append formatting-error instructions to the original complete prompt."""
+        return original_prompt + retry_feedback