Metadata-Version: 2.4
Name: metafit
Version: 0.1.1
Summary: LLM-guided interpretable decision-tree learning
License-Expression: MIT
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: scipy
Requires-Dist: scikit-learn
Requires-Dist: matplotlib
Requires-Dist: tqdm
Requires-Dist: openai
Requires-Dist: tokenizers
Requires-Dist: draccus
Provides-Extra: local
Requires-Dist: torch; extra == "local"
Requires-Dist: transformers; extra == "local"
Dynamic: license-file

# MetaFit

LLM-guided interpretable decision-tree learning. For a worked example, see
[the example notebook](example/example.ipynb). To replicate our experiments, see [README](exp/README.md)  

see official repo: https://github.com/newchat111/MetaFit-An-Interpretable-Classification-Framework-Under-Small-Sample-Regimes

## Installation
```bash
git clone git@github.com:newchat111/MetaFit-An-Interpretable-Classification-Framework-Under-Small-Sample-Regimes.git
```

```bash
pip install metafit
```

## Quick start: use your own dataset

Run the following blocks in order. Replace the CSV paths, target column,
dataset descriptions, and model name with your own values. The CSVs should
contain numeric features with missing values handled in advance, and use the
same columns and label encoding. 

### 1. Load and split the data

```python
import pandas as pd
from metafit.utils.process import stratified_split_indices

output_col = "<target>"  # Your label column.
development_df = pd.read_csv("<development.csv>")
test_df = pd.read_csv("<test.csv>")

# Sample 48 training rows and 32 validation rows, preserving class proportions.
indices = stratified_split_indices(
    development_df[output_col], train_n=48, val_n=32, random_state=0
)
train_df = development_df.iloc[indices["train"]].reset_index(drop=True)
val_df = development_df.iloc[indices["val"]].reset_index(drop=True)
```

### 2. Describe the task (configure metadata)

These descriptions give the LLM context about your data.

```python
metadata = {
    "task": "classification",
    "output_col": output_col,
    "role_description": "You are an expert in the domain of this classification task.",
    "context_description": (
        "<Explain what each row represents and what each target label means.>\n"
        "<Describe each feature, including its units or encoded values.>"
    ),
}
```

### 3. Connect a model and configure training

Set the API key, endpoint URL, and model explicitly below.

```python
import os
from metafit.agents import APIAgent
from metafit.trainer import Trainer


agent = APIAgent(
    api_key="<your-api-key>", #Replace with your api key
    base_url="<provider's API URL>",  # Replace with your provider's API URL.
    model_name="<model-name>", #Replace with the model name (we used deepseek-v4-pro in our experiments)
)

trainer = Trainer(
    agent=agent,
    data_config=metadata,
    train_df=train_df,
    val_df=val_df,
    train_method="exact_tree",  # Use the numeric thresholds proposed by the LLM.
    update_n=3,                 # Number of model updates.
    batch_size=15,              # Labeled examples shown per update.
    return_weights="best",     # Return the highest-scoring candidate.
    output_dir="runs/quickstart",
)
```

To use `interval_tree`, replace the `Trainer(...)` configuration above with
this one. The LLM proposes a tree with threshold ranges. Differential evolution then searches within those ranges for numeric thresholds that maximize the chosen metric on `val_df`.

We recommend using `interval_tree` when the training set has n ≥ 60 and validation set has n ≥ 40.

```python
trainer = Trainer(
    agent=agent,
    data_config=metadata,
    train_df=train_df,
    val_df=val_df,
    train_method="interval_tree",
    de_config={
        "maxiter": 100,  # Maximum number of optimizer iterations.
        "popsize": 15,   # Population-size multiplier for the search.
        "seed": 0,       # Random seed for threshold optimization.
    },
    selection_metric="f1",
    update_n=3,
    batch_size=15,
    return_weights="best",
    output_dir="runs/quickstart_interval",
)
```

`de_config` controls differential evolution. Increasing `maxiter` or `popsize`
allows a larger search and takes more time. Call `trainer.fit()` in step 4 for
either strategy; it returns a tree with numeric thresholds.


### 4. Train, inspect, and evaluate

```python
from metafit.utils.eval import evaluate

# This makes model API calls and may take a few minutes.
tree = trainer.fit()
print(tree.show())
print(tree.rationale)

print(evaluate(tree, test_df, output_col))
```

Optionally save a tree diagram:

```python
fig = tree.show_fig("runs/quickstart/tree")
```

This saves `tree.png`, `tree.svg`, and `tree_editable.svg` in `runs/quickstart/`.

### 5. Inference

```python
features = test_df.drop(columns=[output_col])
predictions = tree.predict(features)
print(predictions)
```

To use a local model, replace the `APIAgent(...)` block above with this
`LocalAgent(...)` configuration before creating `Trainer`. This example uses
[Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B).

Install the local-model dependencies, including Accelerate for automatic device placement:

```bash
python -m pip install "metafit[local]" accelerate
```

```python
from metafit.agents import LocalAgent

agent = LocalAgent(
    model_name="Qwen/Qwen3-30B-A3B",
    temperature=0.7,
    max_new_tokens=8192,  # Maximum output length per response.
    model_kwargs={
        "device_map": "auto",  # Automatically place model layers on available devices.
        "dtype": "auto",       # Use the model's saved weight precision.
    },
    generation_kwargs={"top_p": 0.8, "top_k": 20},
    chat_template_kwargs={"enable_thinking": False},  # Generate without reasoning text.
)
```



## Trainer configuration

Pass these arguments to `Trainer(...)`. Defaults below are the constructor
defaults; the quick start above overrides some of them.

| Argument | Default | Meaning |
| --- | --- | --- |
| `agent` | Required | The model connection, such as an `APIAgent` or `LocalAgent`. |
| `data_config` | `None` | Dataset metadata: the target column, task, and descriptions for the LLM. |
| `train_df` | `None` | Labeled data used for examples and training feedback. |
| `val_df` | `None` | Labeled data used to score candidates and optimize thresholds. |
| `train_method` | `"interval_tree"` | `"exact_tree"` uses the LLM's numeric thresholds. `"interval_tree"` optimizes the LLM's threshold ranges using `de_config`. |
| `update_n` | `9` | Total number of model updates. |
| `batch_size` | `15` | Number of labeled examples shown per update, sampled with replacement. `0` shows no examples. |
| `regularizer_candidate_count` | `5` | Number of tree proposals requested per update; at most this many are fitted from each response. |
| `stratify_batches` | `True` | Balance the classes when sampling examples. |
| `misclassified_top_k` | `10` | Maximum number of misclassified training examples included in the next update's feedback. `0` omits these examples. |
| `updates_per_epoch` | `3` | Number of updates before advancing the seed used to sample batches. |
| `selection_metric` | `"f1"` | Metric used to optimize and rank trees: `"accuracy"`, `"precision"`, `"recall"`, `"f1"`, or `"roc_auc"`. |
| `return_weights` | `"llm_polish"` | `"best"` returns the tree with the highest bootstrap validation mean. `"llm_polish"` asks the LLM to choose an existing fitted tree using candidate scores. |
| `bootstrap_selection_iterations` | `100` | Number of validation resamples used to score each candidate. With `"best"`, `0` uses the original validation score. `"llm_polish"` uses a positive count. |
| `bootstrap_selection_sample_size` | `None` | Rows drawn with replacement per validation resample. `None` uses the validation-set size. |
| `random_state` | `114514` | Seed for sampling training batches and bootstrap validation rows. |
| `de_config` | `None` | Differential-evolution settings for `"interval_tree"`, such as `maxiter`, `popsize`, `tol`, and `seed`. Supply a dictionary when using this method. |
| `max_retries` | `3` | Additional generation attempts when a response produces no usable tree. |
| `output_dir` | `"tmp"` | Folder for training logs, saved candidates, and final selection. |
| `resume` | `True` | Restore completed updates from `output_dir`. `False` starts a fresh run. |
| `logger` | `None` | A custom `TrainerLogger`. `None` creates one using `output_dir` and `resume`. |

## License

MetaFit is licensed under the [MIT License](LICENSE).
