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 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).
@@ -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