tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

make sure to reset the session after each test

Browse Source
This commit is contained in:
Wing Lian
2025-04-30 17:21:25 -04:00
parent f4a0f496a0
commit 0ed4b4c310
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/e2e/integrations/test_llm_compressor.py
View File

@@ -5,6 +5,7 @@ E2E smoke tests for LLMCompressorPlugin integration
from pathlib import Path from pathlib import Path
import pytest import pytest
from llmcompressor import active_session
from axolotl.cli.args import TrainerCliArgs from axolotl.cli.args import TrainerCliArgs
from axolotl.common.datasets import load_datasets from axolotl.common.datasets import load_datasets
@@ -89,9 +90,12 @@ class TestLLMCompressorIntegration:
cli_args = TrainerCliArgs() cli_args = TrainerCliArgs()
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args) dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
train(cfg=cfg, dataset_meta=dataset_meta) try:
check_model_output_exists(temp_dir, cfg) train(cfg=cfg, dataset_meta=dataset_meta)
_check_llmcompressor_model_outputs(temp_dir, save_compressed) check_model_output_exists(temp_dir, cfg)
_check_llmcompressor_model_outputs(temp_dir, save_compressed)
finally:
active_session().reset()
def _check_llmcompressor_model_outputs(temp_dir, save_compressed): def _check_llmcompressor_model_outputs(temp_dir, save_compressed):