barebone-ify the test so we get rid of unneeded processes

This commit is contained in:
Sunny Liu
2025-04-08 14:55:14 -04:00
committed by Sung Ching Liu
parent 7d479348ee
commit fefcbc300d

View File

@@ -9,7 +9,7 @@ import unittest
import transformers import transformers
from axolotl.cli.args import PreprocessCliArgs from axolotl.cli.args import PreprocessCliArgs
from axolotl.cli.preprocess import do_preprocess from axolotl.common.datasets import load_preference_datasets
from axolotl.utils.config import normalize_config, validate_config from axolotl.utils.config import normalize_config, validate_config
from axolotl.utils.dict import DictDefault from axolotl.utils.dict import DictDefault
@@ -49,6 +49,7 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
cfg = DictDefault( cfg = DictDefault(
{ {
"base_model": "HuggingFaceTB/SmolLM2-135M", "base_model": "HuggingFaceTB/SmolLM2-135M",
"strict": False,
"rl": "grpo", "rl": "grpo",
"trl": { "trl": {
"beta": 0.001, "beta": 0.001,
@@ -60,10 +61,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
], # format: '{file_name}.{fn_name}' ], # format: '{file_name}.{fn_name}'
"reward_weights": [1.0], "reward_weights": [1.0],
}, },
"vllm": {
"max_model_len": 800,
"enable_prefix_caching": True,
},
"datasets": [ "datasets": [
{ {
"path": "openai/gsm8k", "path": "openai/gsm8k",
@@ -85,4 +82,4 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
parser = transformers.HfArgumentParser(PreprocessCliArgs) parser = transformers.HfArgumentParser(PreprocessCliArgs)
cli_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True) cli_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)
do_preprocess(cfg, cli_args) load_preference_datasets(cfg=cfg, cli_args=cli_args)