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
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.dict import DictDefault
@@ -49,6 +49,7 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
cfg = DictDefault(
{
"base_model": "HuggingFaceTB/SmolLM2-135M",
"strict": False,
"rl": "grpo",
"trl": {
"beta": 0.001,
@@ -60,10 +61,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
], # format: '{file_name}.{fn_name}'
"reward_weights": [1.0],
},
"vllm": {
"max_model_len": 800,
"enable_prefix_caching": True,
},
"datasets": [
{
"path": "openai/gsm8k",
@@ -85,4 +82,4 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
parser = transformers.HfArgumentParser(PreprocessCliArgs)
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)