metafit 0.1.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- metafit-0.1.1/LICENSE +21 -0
- metafit-0.1.1/PKG-INFO +227 -0
- metafit-0.1.1/README.md +205 -0
- metafit-0.1.1/metafit/__init__.py +1 -0
- metafit-0.1.1/metafit/agents.py +162 -0
- metafit-0.1.1/metafit/logger.py +136 -0
- metafit-0.1.1/metafit/prompts/template.py +306 -0
- metafit-0.1.1/metafit/trainer.py +298 -0
- metafit-0.1.1/metafit/tree.py +220 -0
- metafit-0.1.1/metafit/treediff/__init__.py +5 -0
- metafit-0.1.1/metafit/treediff/tree_diff.py +268 -0
- metafit-0.1.1/metafit/treediff/utils.py +227 -0
- metafit-0.1.1/metafit/utils/analysis.py +211 -0
- metafit-0.1.1/metafit/utils/eval.py +59 -0
- metafit-0.1.1/metafit/utils/process.py +287 -0
- metafit-0.1.1/metafit/utils/token_estimator.py +145 -0
- metafit-0.1.1/metafit/utils/trainer_utils.py +116 -0
- metafit-0.1.1/metafit/utils/tree_plot.py +143 -0
- metafit-0.1.1/metafit.egg-info/PKG-INFO +227 -0
- metafit-0.1.1/metafit.egg-info/SOURCES.txt +23 -0
- metafit-0.1.1/metafit.egg-info/dependency_links.txt +1 -0
- metafit-0.1.1/metafit.egg-info/requires.txt +13 -0
- metafit-0.1.1/metafit.egg-info/top_level.txt +1 -0
- metafit-0.1.1/pyproject.toml +35 -0
- metafit-0.1.1/setup.cfg +4 -0
metafit-0.1.1/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 MetaFit contributors
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
metafit-0.1.1/PKG-INFO
ADDED
|
@@ -0,0 +1,227 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: metafit
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: LLM-guided interpretable decision-tree learning
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Requires-Dist: numpy
|
|
10
|
+
Requires-Dist: pandas
|
|
11
|
+
Requires-Dist: scipy
|
|
12
|
+
Requires-Dist: scikit-learn
|
|
13
|
+
Requires-Dist: matplotlib
|
|
14
|
+
Requires-Dist: tqdm
|
|
15
|
+
Requires-Dist: openai
|
|
16
|
+
Requires-Dist: tokenizers
|
|
17
|
+
Requires-Dist: draccus
|
|
18
|
+
Provides-Extra: local
|
|
19
|
+
Requires-Dist: torch; extra == "local"
|
|
20
|
+
Requires-Dist: transformers; extra == "local"
|
|
21
|
+
Dynamic: license-file
|
|
22
|
+
|
|
23
|
+
# MetaFit
|
|
24
|
+
|
|
25
|
+
LLM-guided interpretable decision-tree learning. For a worked example, see
|
|
26
|
+
[the example notebook](example/example.ipynb). To replicate our experiments, see [README](exp/README.md)
|
|
27
|
+
|
|
28
|
+
see official repo: https://github.com/newchat111/MetaFit-An-Interpretable-Classification-Framework-Under-Small-Sample-Regimes
|
|
29
|
+
|
|
30
|
+
## Installation
|
|
31
|
+
```bash
|
|
32
|
+
git clone git@github.com:newchat111/MetaFit-An-Interpretable-Classification-Framework-Under-Small-Sample-Regimes.git
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
pip install metafit
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
## Quick start: use your own dataset
|
|
40
|
+
|
|
41
|
+
Run the following blocks in order. Replace the CSV paths, target column,
|
|
42
|
+
dataset descriptions, and model name with your own values. The CSVs should
|
|
43
|
+
contain numeric features with missing values handled in advance, and use the
|
|
44
|
+
same columns and label encoding.
|
|
45
|
+
|
|
46
|
+
### 1. Load and split the data
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
import pandas as pd
|
|
50
|
+
from metafit.utils.process import stratified_split_indices
|
|
51
|
+
|
|
52
|
+
output_col = "<target>" # Your label column.
|
|
53
|
+
development_df = pd.read_csv("<development.csv>")
|
|
54
|
+
test_df = pd.read_csv("<test.csv>")
|
|
55
|
+
|
|
56
|
+
# Sample 48 training rows and 32 validation rows, preserving class proportions.
|
|
57
|
+
indices = stratified_split_indices(
|
|
58
|
+
development_df[output_col], train_n=48, val_n=32, random_state=0
|
|
59
|
+
)
|
|
60
|
+
train_df = development_df.iloc[indices["train"]].reset_index(drop=True)
|
|
61
|
+
val_df = development_df.iloc[indices["val"]].reset_index(drop=True)
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
### 2. Describe the task (configure metadata)
|
|
65
|
+
|
|
66
|
+
These descriptions give the LLM context about your data.
|
|
67
|
+
|
|
68
|
+
```python
|
|
69
|
+
metadata = {
|
|
70
|
+
"task": "classification",
|
|
71
|
+
"output_col": output_col,
|
|
72
|
+
"role_description": "You are an expert in the domain of this classification task.",
|
|
73
|
+
"context_description": (
|
|
74
|
+
"<Explain what each row represents and what each target label means.>\n"
|
|
75
|
+
"<Describe each feature, including its units or encoded values.>"
|
|
76
|
+
),
|
|
77
|
+
}
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
### 3. Connect a model and configure training
|
|
81
|
+
|
|
82
|
+
Set the API key, endpoint URL, and model explicitly below.
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
import os
|
|
86
|
+
from metafit.agents import APIAgent
|
|
87
|
+
from metafit.trainer import Trainer
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
agent = APIAgent(
|
|
91
|
+
api_key="<your-api-key>", #Replace with your api key
|
|
92
|
+
base_url="<provider's API URL>", # Replace with your provider's API URL.
|
|
93
|
+
model_name="<model-name>", #Replace with the model name (we used deepseek-v4-pro in our experiments)
|
|
94
|
+
)
|
|
95
|
+
|
|
96
|
+
trainer = Trainer(
|
|
97
|
+
agent=agent,
|
|
98
|
+
data_config=metadata,
|
|
99
|
+
train_df=train_df,
|
|
100
|
+
val_df=val_df,
|
|
101
|
+
train_method="exact_tree", # Use the numeric thresholds proposed by the LLM.
|
|
102
|
+
update_n=3, # Number of model updates.
|
|
103
|
+
batch_size=15, # Labeled examples shown per update.
|
|
104
|
+
return_weights="best", # Return the highest-scoring candidate.
|
|
105
|
+
output_dir="runs/quickstart",
|
|
106
|
+
)
|
|
107
|
+
```
|
|
108
|
+
|
|
109
|
+
To use `interval_tree`, replace the `Trainer(...)` configuration above with
|
|
110
|
+
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`.
|
|
111
|
+
|
|
112
|
+
We recommend using `interval_tree` when the training set has n ≥ 60 and validation set has n ≥ 40.
|
|
113
|
+
|
|
114
|
+
```python
|
|
115
|
+
trainer = Trainer(
|
|
116
|
+
agent=agent,
|
|
117
|
+
data_config=metadata,
|
|
118
|
+
train_df=train_df,
|
|
119
|
+
val_df=val_df,
|
|
120
|
+
train_method="interval_tree",
|
|
121
|
+
de_config={
|
|
122
|
+
"maxiter": 100, # Maximum number of optimizer iterations.
|
|
123
|
+
"popsize": 15, # Population-size multiplier for the search.
|
|
124
|
+
"seed": 0, # Random seed for threshold optimization.
|
|
125
|
+
},
|
|
126
|
+
selection_metric="f1",
|
|
127
|
+
update_n=3,
|
|
128
|
+
batch_size=15,
|
|
129
|
+
return_weights="best",
|
|
130
|
+
output_dir="runs/quickstart_interval",
|
|
131
|
+
)
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
`de_config` controls differential evolution. Increasing `maxiter` or `popsize`
|
|
135
|
+
allows a larger search and takes more time. Call `trainer.fit()` in step 4 for
|
|
136
|
+
either strategy; it returns a tree with numeric thresholds.
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
### 4. Train, inspect, and evaluate
|
|
140
|
+
|
|
141
|
+
```python
|
|
142
|
+
from metafit.utils.eval import evaluate
|
|
143
|
+
|
|
144
|
+
# This makes model API calls and may take a few minutes.
|
|
145
|
+
tree = trainer.fit()
|
|
146
|
+
print(tree.show())
|
|
147
|
+
print(tree.rationale)
|
|
148
|
+
|
|
149
|
+
print(evaluate(tree, test_df, output_col))
|
|
150
|
+
```
|
|
151
|
+
|
|
152
|
+
Optionally save a tree diagram:
|
|
153
|
+
|
|
154
|
+
```python
|
|
155
|
+
fig = tree.show_fig("runs/quickstart/tree")
|
|
156
|
+
```
|
|
157
|
+
|
|
158
|
+
This saves `tree.png`, `tree.svg`, and `tree_editable.svg` in `runs/quickstart/`.
|
|
159
|
+
|
|
160
|
+
### 5. Inference
|
|
161
|
+
|
|
162
|
+
```python
|
|
163
|
+
features = test_df.drop(columns=[output_col])
|
|
164
|
+
predictions = tree.predict(features)
|
|
165
|
+
print(predictions)
|
|
166
|
+
```
|
|
167
|
+
|
|
168
|
+
To use a local model, replace the `APIAgent(...)` block above with this
|
|
169
|
+
`LocalAgent(...)` configuration before creating `Trainer`. This example uses
|
|
170
|
+
[Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B).
|
|
171
|
+
|
|
172
|
+
Install the local-model dependencies, including Accelerate for automatic device placement:
|
|
173
|
+
|
|
174
|
+
```bash
|
|
175
|
+
python -m pip install "metafit[local]" accelerate
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
```python
|
|
179
|
+
from metafit.agents import LocalAgent
|
|
180
|
+
|
|
181
|
+
agent = LocalAgent(
|
|
182
|
+
model_name="Qwen/Qwen3-30B-A3B",
|
|
183
|
+
temperature=0.7,
|
|
184
|
+
max_new_tokens=8192, # Maximum output length per response.
|
|
185
|
+
model_kwargs={
|
|
186
|
+
"device_map": "auto", # Automatically place model layers on available devices.
|
|
187
|
+
"dtype": "auto", # Use the model's saved weight precision.
|
|
188
|
+
},
|
|
189
|
+
generation_kwargs={"top_p": 0.8, "top_k": 20},
|
|
190
|
+
chat_template_kwargs={"enable_thinking": False}, # Generate without reasoning text.
|
|
191
|
+
)
|
|
192
|
+
```
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
## Trainer configuration
|
|
197
|
+
|
|
198
|
+
Pass these arguments to `Trainer(...)`. Defaults below are the constructor
|
|
199
|
+
defaults; the quick start above overrides some of them.
|
|
200
|
+
|
|
201
|
+
| Argument | Default | Meaning |
|
|
202
|
+
| --- | --- | --- |
|
|
203
|
+
| `agent` | Required | The model connection, such as an `APIAgent` or `LocalAgent`. |
|
|
204
|
+
| `data_config` | `None` | Dataset metadata: the target column, task, and descriptions for the LLM. |
|
|
205
|
+
| `train_df` | `None` | Labeled data used for examples and training feedback. |
|
|
206
|
+
| `val_df` | `None` | Labeled data used to score candidates and optimize thresholds. |
|
|
207
|
+
| `train_method` | `"interval_tree"` | `"exact_tree"` uses the LLM's numeric thresholds. `"interval_tree"` optimizes the LLM's threshold ranges using `de_config`. |
|
|
208
|
+
| `update_n` | `9` | Total number of model updates. |
|
|
209
|
+
| `batch_size` | `15` | Number of labeled examples shown per update, sampled with replacement. `0` shows no examples. |
|
|
210
|
+
| `regularizer_candidate_count` | `5` | Number of tree proposals requested per update; at most this many are fitted from each response. |
|
|
211
|
+
| `stratify_batches` | `True` | Balance the classes when sampling examples. |
|
|
212
|
+
| `misclassified_top_k` | `10` | Maximum number of misclassified training examples included in the next update's feedback. `0` omits these examples. |
|
|
213
|
+
| `updates_per_epoch` | `3` | Number of updates before advancing the seed used to sample batches. |
|
|
214
|
+
| `selection_metric` | `"f1"` | Metric used to optimize and rank trees: `"accuracy"`, `"precision"`, `"recall"`, `"f1"`, or `"roc_auc"`. |
|
|
215
|
+
| `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. |
|
|
216
|
+
| `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. |
|
|
217
|
+
| `bootstrap_selection_sample_size` | `None` | Rows drawn with replacement per validation resample. `None` uses the validation-set size. |
|
|
218
|
+
| `random_state` | `114514` | Seed for sampling training batches and bootstrap validation rows. |
|
|
219
|
+
| `de_config` | `None` | Differential-evolution settings for `"interval_tree"`, such as `maxiter`, `popsize`, `tol`, and `seed`. Supply a dictionary when using this method. |
|
|
220
|
+
| `max_retries` | `3` | Additional generation attempts when a response produces no usable tree. |
|
|
221
|
+
| `output_dir` | `"tmp"` | Folder for training logs, saved candidates, and final selection. |
|
|
222
|
+
| `resume` | `True` | Restore completed updates from `output_dir`. `False` starts a fresh run. |
|
|
223
|
+
| `logger` | `None` | A custom `TrainerLogger`. `None` creates one using `output_dir` and `resume`. |
|
|
224
|
+
|
|
225
|
+
## License
|
|
226
|
+
|
|
227
|
+
MetaFit is licensed under the [MIT License](LICENSE).
|
metafit-0.1.1/README.md
ADDED
|
@@ -0,0 +1,205 @@
|
|
|
1
|
+
# MetaFit
|
|
2
|
+
|
|
3
|
+
LLM-guided interpretable decision-tree learning. For a worked example, see
|
|
4
|
+
[the example notebook](example/example.ipynb). To replicate our experiments, see [README](exp/README.md)
|
|
5
|
+
|
|
6
|
+
see official repo: https://github.com/newchat111/MetaFit-An-Interpretable-Classification-Framework-Under-Small-Sample-Regimes
|
|
7
|
+
|
|
8
|
+
## Installation
|
|
9
|
+
```bash
|
|
10
|
+
git clone git@github.com:newchat111/MetaFit-An-Interpretable-Classification-Framework-Under-Small-Sample-Regimes.git
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
```bash
|
|
14
|
+
pip install metafit
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
## Quick start: use your own dataset
|
|
18
|
+
|
|
19
|
+
Run the following blocks in order. Replace the CSV paths, target column,
|
|
20
|
+
dataset descriptions, and model name with your own values. The CSVs should
|
|
21
|
+
contain numeric features with missing values handled in advance, and use the
|
|
22
|
+
same columns and label encoding.
|
|
23
|
+
|
|
24
|
+
### 1. Load and split the data
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
import pandas as pd
|
|
28
|
+
from metafit.utils.process import stratified_split_indices
|
|
29
|
+
|
|
30
|
+
output_col = "<target>" # Your label column.
|
|
31
|
+
development_df = pd.read_csv("<development.csv>")
|
|
32
|
+
test_df = pd.read_csv("<test.csv>")
|
|
33
|
+
|
|
34
|
+
# Sample 48 training rows and 32 validation rows, preserving class proportions.
|
|
35
|
+
indices = stratified_split_indices(
|
|
36
|
+
development_df[output_col], train_n=48, val_n=32, random_state=0
|
|
37
|
+
)
|
|
38
|
+
train_df = development_df.iloc[indices["train"]].reset_index(drop=True)
|
|
39
|
+
val_df = development_df.iloc[indices["val"]].reset_index(drop=True)
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
### 2. Describe the task (configure metadata)
|
|
43
|
+
|
|
44
|
+
These descriptions give the LLM context about your data.
|
|
45
|
+
|
|
46
|
+
```python
|
|
47
|
+
metadata = {
|
|
48
|
+
"task": "classification",
|
|
49
|
+
"output_col": output_col,
|
|
50
|
+
"role_description": "You are an expert in the domain of this classification task.",
|
|
51
|
+
"context_description": (
|
|
52
|
+
"<Explain what each row represents and what each target label means.>\n"
|
|
53
|
+
"<Describe each feature, including its units or encoded values.>"
|
|
54
|
+
),
|
|
55
|
+
}
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
### 3. Connect a model and configure training
|
|
59
|
+
|
|
60
|
+
Set the API key, endpoint URL, and model explicitly below.
|
|
61
|
+
|
|
62
|
+
```python
|
|
63
|
+
import os
|
|
64
|
+
from metafit.agents import APIAgent
|
|
65
|
+
from metafit.trainer import Trainer
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
agent = APIAgent(
|
|
69
|
+
api_key="<your-api-key>", #Replace with your api key
|
|
70
|
+
base_url="<provider's API URL>", # Replace with your provider's API URL.
|
|
71
|
+
model_name="<model-name>", #Replace with the model name (we used deepseek-v4-pro in our experiments)
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
trainer = Trainer(
|
|
75
|
+
agent=agent,
|
|
76
|
+
data_config=metadata,
|
|
77
|
+
train_df=train_df,
|
|
78
|
+
val_df=val_df,
|
|
79
|
+
train_method="exact_tree", # Use the numeric thresholds proposed by the LLM.
|
|
80
|
+
update_n=3, # Number of model updates.
|
|
81
|
+
batch_size=15, # Labeled examples shown per update.
|
|
82
|
+
return_weights="best", # Return the highest-scoring candidate.
|
|
83
|
+
output_dir="runs/quickstart",
|
|
84
|
+
)
|
|
85
|
+
```
|
|
86
|
+
|
|
87
|
+
To use `interval_tree`, replace the `Trainer(...)` configuration above with
|
|
88
|
+
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`.
|
|
89
|
+
|
|
90
|
+
We recommend using `interval_tree` when the training set has n ≥ 60 and validation set has n ≥ 40.
|
|
91
|
+
|
|
92
|
+
```python
|
|
93
|
+
trainer = Trainer(
|
|
94
|
+
agent=agent,
|
|
95
|
+
data_config=metadata,
|
|
96
|
+
train_df=train_df,
|
|
97
|
+
val_df=val_df,
|
|
98
|
+
train_method="interval_tree",
|
|
99
|
+
de_config={
|
|
100
|
+
"maxiter": 100, # Maximum number of optimizer iterations.
|
|
101
|
+
"popsize": 15, # Population-size multiplier for the search.
|
|
102
|
+
"seed": 0, # Random seed for threshold optimization.
|
|
103
|
+
},
|
|
104
|
+
selection_metric="f1",
|
|
105
|
+
update_n=3,
|
|
106
|
+
batch_size=15,
|
|
107
|
+
return_weights="best",
|
|
108
|
+
output_dir="runs/quickstart_interval",
|
|
109
|
+
)
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
`de_config` controls differential evolution. Increasing `maxiter` or `popsize`
|
|
113
|
+
allows a larger search and takes more time. Call `trainer.fit()` in step 4 for
|
|
114
|
+
either strategy; it returns a tree with numeric thresholds.
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
### 4. Train, inspect, and evaluate
|
|
118
|
+
|
|
119
|
+
```python
|
|
120
|
+
from metafit.utils.eval import evaluate
|
|
121
|
+
|
|
122
|
+
# This makes model API calls and may take a few minutes.
|
|
123
|
+
tree = trainer.fit()
|
|
124
|
+
print(tree.show())
|
|
125
|
+
print(tree.rationale)
|
|
126
|
+
|
|
127
|
+
print(evaluate(tree, test_df, output_col))
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
Optionally save a tree diagram:
|
|
131
|
+
|
|
132
|
+
```python
|
|
133
|
+
fig = tree.show_fig("runs/quickstart/tree")
|
|
134
|
+
```
|
|
135
|
+
|
|
136
|
+
This saves `tree.png`, `tree.svg`, and `tree_editable.svg` in `runs/quickstart/`.
|
|
137
|
+
|
|
138
|
+
### 5. Inference
|
|
139
|
+
|
|
140
|
+
```python
|
|
141
|
+
features = test_df.drop(columns=[output_col])
|
|
142
|
+
predictions = tree.predict(features)
|
|
143
|
+
print(predictions)
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
To use a local model, replace the `APIAgent(...)` block above with this
|
|
147
|
+
`LocalAgent(...)` configuration before creating `Trainer`. This example uses
|
|
148
|
+
[Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B).
|
|
149
|
+
|
|
150
|
+
Install the local-model dependencies, including Accelerate for automatic device placement:
|
|
151
|
+
|
|
152
|
+
```bash
|
|
153
|
+
python -m pip install "metafit[local]" accelerate
|
|
154
|
+
```
|
|
155
|
+
|
|
156
|
+
```python
|
|
157
|
+
from metafit.agents import LocalAgent
|
|
158
|
+
|
|
159
|
+
agent = LocalAgent(
|
|
160
|
+
model_name="Qwen/Qwen3-30B-A3B",
|
|
161
|
+
temperature=0.7,
|
|
162
|
+
max_new_tokens=8192, # Maximum output length per response.
|
|
163
|
+
model_kwargs={
|
|
164
|
+
"device_map": "auto", # Automatically place model layers on available devices.
|
|
165
|
+
"dtype": "auto", # Use the model's saved weight precision.
|
|
166
|
+
},
|
|
167
|
+
generation_kwargs={"top_p": 0.8, "top_k": 20},
|
|
168
|
+
chat_template_kwargs={"enable_thinking": False}, # Generate without reasoning text.
|
|
169
|
+
)
|
|
170
|
+
```
|
|
171
|
+
|
|
172
|
+
|
|
173
|
+
|
|
174
|
+
## Trainer configuration
|
|
175
|
+
|
|
176
|
+
Pass these arguments to `Trainer(...)`. Defaults below are the constructor
|
|
177
|
+
defaults; the quick start above overrides some of them.
|
|
178
|
+
|
|
179
|
+
| Argument | Default | Meaning |
|
|
180
|
+
| --- | --- | --- |
|
|
181
|
+
| `agent` | Required | The model connection, such as an `APIAgent` or `LocalAgent`. |
|
|
182
|
+
| `data_config` | `None` | Dataset metadata: the target column, task, and descriptions for the LLM. |
|
|
183
|
+
| `train_df` | `None` | Labeled data used for examples and training feedback. |
|
|
184
|
+
| `val_df` | `None` | Labeled data used to score candidates and optimize thresholds. |
|
|
185
|
+
| `train_method` | `"interval_tree"` | `"exact_tree"` uses the LLM's numeric thresholds. `"interval_tree"` optimizes the LLM's threshold ranges using `de_config`. |
|
|
186
|
+
| `update_n` | `9` | Total number of model updates. |
|
|
187
|
+
| `batch_size` | `15` | Number of labeled examples shown per update, sampled with replacement. `0` shows no examples. |
|
|
188
|
+
| `regularizer_candidate_count` | `5` | Number of tree proposals requested per update; at most this many are fitted from each response. |
|
|
189
|
+
| `stratify_batches` | `True` | Balance the classes when sampling examples. |
|
|
190
|
+
| `misclassified_top_k` | `10` | Maximum number of misclassified training examples included in the next update's feedback. `0` omits these examples. |
|
|
191
|
+
| `updates_per_epoch` | `3` | Number of updates before advancing the seed used to sample batches. |
|
|
192
|
+
| `selection_metric` | `"f1"` | Metric used to optimize and rank trees: `"accuracy"`, `"precision"`, `"recall"`, `"f1"`, or `"roc_auc"`. |
|
|
193
|
+
| `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. |
|
|
194
|
+
| `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. |
|
|
195
|
+
| `bootstrap_selection_sample_size` | `None` | Rows drawn with replacement per validation resample. `None` uses the validation-set size. |
|
|
196
|
+
| `random_state` | `114514` | Seed for sampling training batches and bootstrap validation rows. |
|
|
197
|
+
| `de_config` | `None` | Differential-evolution settings for `"interval_tree"`, such as `maxiter`, `popsize`, `tol`, and `seed`. Supply a dictionary when using this method. |
|
|
198
|
+
| `max_retries` | `3` | Additional generation attempts when a response produces no usable tree. |
|
|
199
|
+
| `output_dir` | `"tmp"` | Folder for training logs, saved candidates, and final selection. |
|
|
200
|
+
| `resume` | `True` | Restore completed updates from `output_dir`. `False` starts a fresh run. |
|
|
201
|
+
| `logger` | `None` | A custom `TrainerLogger`. `None` creates one using `output_dir` and `resume`. |
|
|
202
|
+
|
|
203
|
+
## License
|
|
204
|
+
|
|
205
|
+
MetaFit is licensed under the [MIT License](LICENSE).
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""MetaFit rework: agents, decision trees, and training."""
|
|
@@ -0,0 +1,162 @@
|
|
|
1
|
+
"""Text generation with local Hugging Face models or OpenAI-compatible APIs.
|
|
2
|
+
|
|
3
|
+
Dependencies are imported only when their respective agent is initialized.
|
|
4
|
+
Each call to ``generate`` is independent and returns the generated text.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from abc import ABC, abstractmethod
|
|
8
|
+
from typing import Any
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class Agent(ABC):
|
|
12
|
+
"""Shared interface for language-model agents."""
|
|
13
|
+
|
|
14
|
+
@abstractmethod
|
|
15
|
+
def generate(
|
|
16
|
+
self,
|
|
17
|
+
prompt: str,
|
|
18
|
+
system_prompt: str | None = None,
|
|
19
|
+
**generation_kwargs: Any,
|
|
20
|
+
) -> str:
|
|
21
|
+
"""Return generated text for a prompt."""
|
|
22
|
+
pass
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class LocalAgent(Agent):
|
|
26
|
+
"""Load a causal model with separate loading and generation options.
|
|
27
|
+
|
|
28
|
+
``model_kwargs`` accepts dtype, device_map, quantization_config, revision,
|
|
29
|
+
and other model-loading options. ``tokenizer_kwargs`` configures tokenizer
|
|
30
|
+
loading. ``generation_kwargs`` and ``chat_template_kwargs`` store defaults
|
|
31
|
+
that can be overridden in ``generate``.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
def __init__(
|
|
35
|
+
self,
|
|
36
|
+
model_name: str,
|
|
37
|
+
temperature: float = 0.7,
|
|
38
|
+
*,
|
|
39
|
+
device: str | None = None,
|
|
40
|
+
max_new_tokens: int = 512,
|
|
41
|
+
model_kwargs: dict[str, Any] | None = None,
|
|
42
|
+
tokenizer_kwargs: dict[str, Any] | None = None,
|
|
43
|
+
generation_kwargs: dict[str, Any] | None = None,
|
|
44
|
+
chat_template_kwargs: dict[str, Any] | None = None,
|
|
45
|
+
) -> None:
|
|
46
|
+
import torch
|
|
47
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
48
|
+
|
|
49
|
+
self.model_name = model_name
|
|
50
|
+
self.temperature = temperature
|
|
51
|
+
self.max_new_tokens = max_new_tokens
|
|
52
|
+
self.generation_kwargs = dict(generation_kwargs or {})
|
|
53
|
+
self.chat_template_kwargs = dict(chat_template_kwargs or {})
|
|
54
|
+
self.tokenizer = AutoTokenizer.from_pretrained(model_name, **(tokenizer_kwargs or {}))
|
|
55
|
+
self.model = AutoModelForCausalLM.from_pretrained(model_name, **(model_kwargs or {}))
|
|
56
|
+
if (model_kwargs or {}).get("device_map") is None:
|
|
57
|
+
self.model.to(device or ("cuda" if torch.cuda.is_available() else "cpu"))
|
|
58
|
+
self.model.eval()
|
|
59
|
+
|
|
60
|
+
def generate(
|
|
61
|
+
self,
|
|
62
|
+
prompt: str,
|
|
63
|
+
system_prompt: str | None = None,
|
|
64
|
+
*,
|
|
65
|
+
chat_template_kwargs: dict[str, Any] | None = None,
|
|
66
|
+
**generation_kwargs: Any,
|
|
67
|
+
) -> str:
|
|
68
|
+
"""Generate text, overriding stored generation and template defaults.
|
|
69
|
+
|
|
70
|
+
A zero temperature selects greedy decoding. Models with a chat template
|
|
71
|
+
receive formatted messages; other models receive plain text.
|
|
72
|
+
"""
|
|
73
|
+
import torch
|
|
74
|
+
|
|
75
|
+
if self.tokenizer.chat_template:
|
|
76
|
+
messages = []
|
|
77
|
+
if system_prompt:
|
|
78
|
+
messages.append({"role": "system", "content": system_prompt})
|
|
79
|
+
messages.append({"role": "user", "content": prompt})
|
|
80
|
+
text = self.tokenizer.apply_chat_template(
|
|
81
|
+
messages,
|
|
82
|
+
tokenize=False,
|
|
83
|
+
add_generation_prompt=True,
|
|
84
|
+
**{**self.chat_template_kwargs, **(chat_template_kwargs or {})},
|
|
85
|
+
)
|
|
86
|
+
inputs = self.tokenizer(
|
|
87
|
+
text, return_tensors="pt", add_special_tokens=False
|
|
88
|
+
)
|
|
89
|
+
else:
|
|
90
|
+
text = f"{system_prompt}\n\n{prompt}" if system_prompt else prompt
|
|
91
|
+
inputs = self.tokenizer(text, return_tensors="pt")
|
|
92
|
+
inputs = inputs.to(self.model.device)
|
|
93
|
+
|
|
94
|
+
params = {
|
|
95
|
+
"max_new_tokens": self.max_new_tokens,
|
|
96
|
+
"temperature": self.temperature,
|
|
97
|
+
**self.generation_kwargs,
|
|
98
|
+
**generation_kwargs,
|
|
99
|
+
}
|
|
100
|
+
temperature = params.pop("temperature")
|
|
101
|
+
params.setdefault("do_sample", temperature > 0)
|
|
102
|
+
if params["do_sample"]:
|
|
103
|
+
params["temperature"] = temperature
|
|
104
|
+
pad_token_id = self.tokenizer.pad_token_id
|
|
105
|
+
if pad_token_id is None:
|
|
106
|
+
pad_token_id = self.tokenizer.eos_token_id
|
|
107
|
+
if pad_token_id is not None:
|
|
108
|
+
params.setdefault("pad_token_id", pad_token_id)
|
|
109
|
+
|
|
110
|
+
with torch.inference_mode():
|
|
111
|
+
output = self.model.generate(**inputs, **params)
|
|
112
|
+
sequences = output.sequences if hasattr(output, "sequences") else output
|
|
113
|
+
new_tokens = sequences[0, inputs["input_ids"].shape[-1]:]
|
|
114
|
+
return self.tokenizer.decode(new_tokens, skip_special_tokens=True)
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
class APIAgent(Agent):
|
|
118
|
+
"""Generate text with OpenAI Chat Completions or a compatible endpoint.
|
|
119
|
+
|
|
120
|
+
``client_kwargs`` accepts timeout, max_retries, and other client options.
|
|
121
|
+
``generation_kwargs`` stores endpoint-specific defaults such as top_p,
|
|
122
|
+
max_completion_tokens, response_format, reasoning_effort, and extra_body.
|
|
123
|
+
Use ``temperature=None`` to omit temperature from requests.
|
|
124
|
+
"""
|
|
125
|
+
|
|
126
|
+
def __init__(
|
|
127
|
+
self,
|
|
128
|
+
api_key: str | None = None,
|
|
129
|
+
model_name: str = "gpt-4",
|
|
130
|
+
temperature: float | None = 0.7,
|
|
131
|
+
*,
|
|
132
|
+
base_url: str | None = None,
|
|
133
|
+
client_kwargs: dict[str, Any] | None = None,
|
|
134
|
+
generation_kwargs: dict[str, Any] | None = None,
|
|
135
|
+
) -> None:
|
|
136
|
+
"""Use the supplied key, or OPENAI_API_KEY when it is omitted."""
|
|
137
|
+
from openai import OpenAI
|
|
138
|
+
|
|
139
|
+
self.model_name = model_name
|
|
140
|
+
self.temperature = temperature
|
|
141
|
+
self.generation_kwargs = dict(generation_kwargs or {})
|
|
142
|
+
self.client = OpenAI(api_key=api_key, base_url=base_url, **(client_kwargs or {}))
|
|
143
|
+
|
|
144
|
+
def generate(
|
|
145
|
+
self,
|
|
146
|
+
prompt: str,
|
|
147
|
+
system_prompt: str | None = None,
|
|
148
|
+
**generation_kwargs: Any,
|
|
149
|
+
) -> str:
|
|
150
|
+
"""Return one non-streaming reply, overriding stored generation defaults."""
|
|
151
|
+
messages = []
|
|
152
|
+
if system_prompt:
|
|
153
|
+
messages.append({"role": "system", "content": system_prompt})
|
|
154
|
+
messages.append({"role": "user", "content": prompt})
|
|
155
|
+
|
|
156
|
+
params = {"temperature": self.temperature, **self.generation_kwargs, **generation_kwargs}
|
|
157
|
+
if params["temperature"] is None:
|
|
158
|
+
params.pop("temperature")
|
|
159
|
+
response = self.client.chat.completions.create(
|
|
160
|
+
model=self.model_name, messages=messages, **params
|
|
161
|
+
)
|
|
162
|
+
return response.choices[0].message.content
|