mirror of https://github.com/razor-ai/soup.git
889 lines
34 KiB
Python
889 lines
34 KiB
Python
"""Tests for pretrain task — config, plaintext format, template, routing, sweep."""
|
|
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch as mock_patch
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from soup_cli.config.schema import TEMPLATES, SoupConfig
|
|
|
|
# ─── Config Tests ───────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestPretrainConfig:
|
|
"""Test pretrain task config validation."""
|
|
|
|
def test_pretrain_task_accepted(self):
|
|
"""pretrain task should be a valid task type."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.task == "pretrain"
|
|
|
|
def test_pretrain_default_config(self):
|
|
"""pretrain task should use default training config values."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.epochs == 3
|
|
assert cfg.training.lr == pytest.approx(2e-5)
|
|
|
|
def test_pretrain_with_plaintext_format(self):
|
|
"""pretrain task with plaintext format should validate correctly."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./corpus.txt", "format": "plaintext"},
|
|
)
|
|
assert cfg.data.format == "plaintext"
|
|
|
|
def test_plaintext_format_accepted(self):
|
|
"""plaintext should be a valid data format."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="sft",
|
|
data={"train": "./data.jsonl", "format": "plaintext"},
|
|
)
|
|
assert cfg.data.format == "plaintext"
|
|
|
|
def test_pretrain_full_config(self):
|
|
"""Full pretrain config should validate correctly."""
|
|
cfg = SoupConfig(
|
|
base="meta-llama/Llama-3.1-8B",
|
|
task="pretrain",
|
|
data={"train": "./corpus.jsonl", "format": "plaintext", "max_length": 4096},
|
|
training={
|
|
"epochs": 1,
|
|
"lr": 1e-5,
|
|
"gradient_accumulation_steps": 8,
|
|
"lora": {"r": 64, "alpha": 16},
|
|
"quantization": "4bit",
|
|
},
|
|
)
|
|
assert cfg.task == "pretrain"
|
|
assert cfg.data.max_length == 4096
|
|
assert cfg.training.gradient_accumulation_steps == 8
|
|
|
|
def test_pretrain_unsloth_backend(self):
|
|
"""pretrain task with unsloth backend should validate correctly."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
backend="unsloth",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.backend == "unsloth"
|
|
assert cfg.task == "pretrain"
|
|
|
|
|
|
# ─── MoE Config Tests ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestMoEConfig:
|
|
"""Test MoE-specific config fields."""
|
|
|
|
def test_moe_lora_default_false(self):
|
|
"""moe_lora should default to False."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.moe_lora is False
|
|
|
|
def test_moe_lora_enabled(self):
|
|
"""moe_lora should be settable to True."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"moe_lora": True},
|
|
)
|
|
assert cfg.training.moe_lora is True
|
|
|
|
def test_moe_aux_loss_coeff_default(self):
|
|
"""moe_aux_loss_coeff should default to 0.01."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.moe_aux_loss_coeff == pytest.approx(0.01)
|
|
|
|
def test_moe_aux_loss_coeff_custom(self):
|
|
"""Custom moe_aux_loss_coeff should be accepted."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"moe_aux_loss_coeff": 0.05},
|
|
)
|
|
assert cfg.training.moe_aux_loss_coeff == pytest.approx(0.05)
|
|
|
|
def test_moe_aux_loss_coeff_zero_allowed(self):
|
|
"""moe_aux_loss_coeff=0 should be valid (disables aux loss)."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"moe_aux_loss_coeff": 0.0},
|
|
)
|
|
assert cfg.training.moe_aux_loss_coeff == 0.0
|
|
|
|
def test_moe_aux_loss_coeff_negative_rejected(self):
|
|
"""Negative moe_aux_loss_coeff should be rejected."""
|
|
with pytest.raises(ValidationError):
|
|
SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"moe_aux_loss_coeff": -0.1},
|
|
)
|
|
|
|
|
|
# ─── Plaintext Data Format Tests ──────────────────────────────────────────
|
|
|
|
|
|
class TestPlaintextDataFormat:
|
|
"""Test plaintext data format detection and conversion."""
|
|
|
|
def test_format_signature_exists(self):
|
|
"""plaintext format signature should be registered."""
|
|
from soup_cli.data.formats import FORMAT_SIGNATURES
|
|
|
|
assert "plaintext" in FORMAT_SIGNATURES
|
|
assert FORMAT_SIGNATURES["plaintext"] == {"text"}
|
|
|
|
def test_detect_plaintext_format(self):
|
|
"""Should auto-detect plaintext format from data keys."""
|
|
from soup_cli.data.formats import detect_format
|
|
|
|
data = [{"text": "This is a document."}]
|
|
assert detect_format(data) == "plaintext"
|
|
|
|
def test_detect_plaintext_with_extra_keys(self):
|
|
"""Should detect plaintext format even with extra keys."""
|
|
from soup_cli.data.formats import detect_format
|
|
|
|
data = [{"text": "Document content.", "id": 1, "source": "web"}]
|
|
assert detect_format(data) == "plaintext"
|
|
|
|
def test_convert_plaintext(self):
|
|
"""Should convert plaintext row correctly."""
|
|
from soup_cli.data.formats import format_to_messages
|
|
|
|
row = {"text": "This is raw text for pre-training."}
|
|
result = format_to_messages(row, "plaintext")
|
|
assert result["text"] == "This is raw text for pre-training."
|
|
|
|
def test_convert_plaintext_preserves_text_only(self):
|
|
"""Output should have exactly the 'text' key."""
|
|
from soup_cli.data.formats import format_to_messages
|
|
|
|
row = {"text": "content", "id": 42, "source": "web"}
|
|
result = format_to_messages(row, "plaintext")
|
|
assert set(result.keys()) == {"text"}
|
|
|
|
def test_convert_plaintext_empty_text_returns_none(self):
|
|
"""Empty text should cause conversion to return None."""
|
|
from soup_cli.data.formats import format_to_messages
|
|
|
|
row = {"text": ""}
|
|
result = format_to_messages(row, "plaintext")
|
|
assert result is None
|
|
|
|
def test_convert_plaintext_whitespace_only_returns_none(self):
|
|
"""Whitespace-only text should cause conversion to return None."""
|
|
from soup_cli.data.formats import format_to_messages
|
|
|
|
row = {"text": " \n "}
|
|
result = format_to_messages(row, "plaintext")
|
|
assert result is None
|
|
|
|
def test_convert_plaintext_missing_text_returns_none(self):
|
|
"""Row missing 'text' key should return None."""
|
|
from soup_cli.data.formats import format_to_messages
|
|
|
|
row = {"content": "some text"}
|
|
result = format_to_messages(row, "plaintext")
|
|
assert result is None
|
|
|
|
def test_plaintext_not_confused_with_chatml(self):
|
|
"""Data with only 'text' key should not be detected as chatml."""
|
|
from soup_cli.data.formats import detect_format
|
|
|
|
data = [{"text": "plain content"}]
|
|
assert detect_format(data) == "plaintext"
|
|
|
|
def test_chatml_not_detected_as_plaintext(self):
|
|
"""Data with 'messages' key should be detected as chatml, not plaintext."""
|
|
from soup_cli.data.formats import detect_format
|
|
|
|
data = [{"messages": [{"role": "user", "content": "hello"}]}]
|
|
assert detect_format(data) == "chatml"
|
|
|
|
|
|
# ─── Text File Loading Tests ──────────────────────────────────────────────
|
|
|
|
|
|
class TestTxtFileLoading:
|
|
"""Test .txt file loading."""
|
|
|
|
def test_load_txt_file(self, tmp_path):
|
|
"""Should load .txt file as list of {text: ...} dicts."""
|
|
from soup_cli.data.loader import load_raw_data
|
|
|
|
txt_file = tmp_path / "corpus.txt"
|
|
txt_file.write_text("Line one\nLine two\nLine three\n", encoding="utf-8")
|
|
data = load_raw_data(txt_file)
|
|
assert len(data) == 3
|
|
assert data[0] == {"text": "Line one"}
|
|
assert data[1] == {"text": "Line two"}
|
|
assert data[2] == {"text": "Line three"}
|
|
|
|
def test_load_txt_skips_empty_lines(self, tmp_path):
|
|
"""Empty lines should be skipped."""
|
|
from soup_cli.data.loader import load_raw_data
|
|
|
|
txt_file = tmp_path / "corpus.txt"
|
|
txt_file.write_text("Line one\n\n\nLine two\n", encoding="utf-8")
|
|
data = load_raw_data(txt_file)
|
|
assert len(data) == 2
|
|
assert data[0] == {"text": "Line one"}
|
|
assert data[1] == {"text": "Line two"}
|
|
|
|
def test_load_txt_empty_file(self, tmp_path):
|
|
"""Empty .txt file should return empty list."""
|
|
from soup_cli.data.loader import load_raw_data
|
|
|
|
txt_file = tmp_path / "empty.txt"
|
|
txt_file.write_text("", encoding="utf-8")
|
|
data = load_raw_data(txt_file)
|
|
assert data == []
|
|
|
|
def test_txt_extension_supported(self):
|
|
"""'.txt' should be in SUPPORTED_EXTENSIONS."""
|
|
from soup_cli.data.loader import SUPPORTED_EXTENSIONS
|
|
|
|
assert ".txt" in SUPPORTED_EXTENSIONS
|
|
|
|
def test_load_txt_strips_whitespace(self, tmp_path):
|
|
"""Lines should have leading/trailing whitespace stripped."""
|
|
from soup_cli.data.loader import load_raw_data
|
|
|
|
txt_file = tmp_path / "corpus.txt"
|
|
txt_file.write_text(" hello \n world \n", encoding="utf-8")
|
|
data = load_raw_data(txt_file)
|
|
assert data[0] == {"text": "hello"}
|
|
assert data[1] == {"text": "world"}
|
|
|
|
|
|
# ─── Template Tests ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestPretrainTemplate:
|
|
"""Test the pretrain template."""
|
|
|
|
def test_pretrain_template_exists(self):
|
|
assert "pretrain" in TEMPLATES
|
|
|
|
def test_pretrain_template_valid_yaml(self):
|
|
import yaml
|
|
|
|
config = yaml.safe_load(TEMPLATES["pretrain"])
|
|
assert config["task"] == "pretrain"
|
|
assert config["data"]["format"] == "plaintext"
|
|
|
|
def test_pretrain_template_valid_config(self):
|
|
import yaml
|
|
|
|
raw = yaml.safe_load(TEMPLATES["pretrain"])
|
|
cfg = SoupConfig(**raw)
|
|
assert cfg.task == "pretrain"
|
|
assert cfg.data.format == "plaintext"
|
|
|
|
def test_moe_template_exists(self):
|
|
assert "moe" in TEMPLATES
|
|
|
|
def test_moe_template_valid_yaml(self):
|
|
import yaml
|
|
|
|
config = yaml.safe_load(TEMPLATES["moe"])
|
|
assert config["training"]["moe_lora"] is True
|
|
assert config["training"]["moe_aux_loss_coeff"] == 0.01
|
|
|
|
def test_moe_template_valid_config(self):
|
|
import yaml
|
|
|
|
raw = yaml.safe_load(TEMPLATES["moe"])
|
|
cfg = SoupConfig(**raw)
|
|
assert cfg.training.moe_lora is True
|
|
assert cfg.training.moe_aux_loss_coeff == pytest.approx(0.01)
|
|
|
|
|
|
# ─── Train Command Routing Tests ──────────────────────────────────────────
|
|
|
|
|
|
class TestPretrainTrainRouting:
|
|
"""Test that train command routes to pretrain trainer."""
|
|
|
|
def test_pretrain_import_exists(self):
|
|
"""PretrainTrainerWrapper should be importable."""
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
assert PretrainTrainerWrapper is not None
|
|
|
|
def test_pretrain_wrapper_init(self):
|
|
"""PretrainTrainerWrapper should initialize without error."""
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = PretrainTrainerWrapper(cfg, device="cpu")
|
|
assert wrapper.config.task == "pretrain"
|
|
assert wrapper.device == "cpu"
|
|
assert wrapper.model is None
|
|
assert wrapper.trainer is None
|
|
|
|
def test_pretrain_wrapper_init_with_options(self):
|
|
"""PretrainTrainerWrapper should accept all constructor options."""
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = PretrainTrainerWrapper(
|
|
cfg, device="cuda", report_to="wandb", deepspeed_config="ds.json",
|
|
)
|
|
assert wrapper.report_to == "wandb"
|
|
assert wrapper.deepspeed_config == "ds.json"
|
|
|
|
|
|
# ─── Sweep Shortcut Tests ─────────────────────────────────────────────────
|
|
|
|
|
|
class TestPretrainSweepParams:
|
|
"""Test pretrain/MoE parameter shortcuts in sweep."""
|
|
|
|
def test_moe_lora_shortcut(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {"training": {"moe_lora": False}}
|
|
_set_nested_param(config, "moe_lora", True)
|
|
assert config["training"]["moe_lora"] is True
|
|
|
|
def test_moe_aux_loss_coeff_shortcut(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {"training": {"moe_aux_loss_coeff": 0.01}}
|
|
_set_nested_param(config, "moe_aux_loss_coeff", 0.05)
|
|
assert config["training"]["moe_aux_loss_coeff"] == pytest.approx(0.05)
|
|
|
|
def test_moe_lora_shortcut_creates_nested_key(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {}
|
|
_set_nested_param(config, "moe_lora", True)
|
|
assert config["training"]["moe_lora"] is True
|
|
|
|
def test_sweep_run_single_routes_to_pretrain_trainer(self):
|
|
"""_run_single should instantiate PretrainTrainerWrapper for pretrain task."""
|
|
from soup_cli.commands.sweep import _run_single
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
|
|
fake_dataset = {
|
|
"train": [{"text": "Some pre-training text."}]
|
|
}
|
|
fake_result = {
|
|
"initial_loss": 3.0,
|
|
"final_loss": 2.5,
|
|
"total_steps": 10,
|
|
"duration_secs": 120.0,
|
|
"output_dir": "./output",
|
|
"duration": "2m",
|
|
}
|
|
|
|
fake_gpu_info = {"memory_total": "0 MB", "memory_total_bytes": 0}
|
|
with mock_patch("soup_cli.data.loader.load_dataset", return_value=fake_dataset), \
|
|
mock_patch("soup_cli.utils.gpu.detect_device", return_value=("cpu", "CPU")), \
|
|
mock_patch("soup_cli.utils.gpu.get_gpu_info", return_value=fake_gpu_info), \
|
|
mock_patch("soup_cli.experiment.tracker.ExperimentTracker") as mock_tracker_cls, \
|
|
mock_patch("soup_cli.monitoring.display.TrainingDisplay"), \
|
|
mock_patch("soup_cli.trainer.pretrain.PretrainTrainerWrapper.setup"), \
|
|
mock_patch(
|
|
"soup_cli.trainer.pretrain.PretrainTrainerWrapper.train",
|
|
return_value=fake_result,
|
|
) as mock_train:
|
|
mock_tracker = MagicMock()
|
|
mock_tracker.start_run.return_value = "run-pretrain-1"
|
|
mock_tracker_cls.return_value = mock_tracker
|
|
|
|
result = _run_single(cfg, {}, "pretrain_run_1", None)
|
|
|
|
mock_train.assert_called_once()
|
|
assert result["run_id"] == "run-pretrain-1"
|
|
|
|
|
|
# ─── Train Guard Test ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestPretrainTrainGuard:
|
|
"""Test the RuntimeError guard when train() is called before setup()."""
|
|
|
|
def test_train_before_setup_raises_runtime_error(self):
|
|
"""Calling train() before setup() should raise RuntimeError."""
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = PretrainTrainerWrapper(cfg)
|
|
with pytest.raises(RuntimeError, match="setup\\(dataset\\) first"):
|
|
wrapper.train()
|
|
|
|
def test_train_error_message_mentions_setup(self):
|
|
"""RuntimeError message should mention setup()."""
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = PretrainTrainerWrapper(cfg)
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
wrapper.train()
|
|
assert "setup" in str(exc_info.value).lower()
|
|
|
|
|
|
# ─── Train Method Result Structure ──────────────────────────────────────────
|
|
|
|
|
|
class TestPretrainTrainResults:
|
|
"""Test the result dict returned by train() using a mocked trainer."""
|
|
|
|
def _make_wrapper_with_mock_trainer(self, log_history=None, global_step=20):
|
|
"""Helper: return a PretrainTrainerWrapper with trainer pre-injected."""
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
output="./output",
|
|
)
|
|
wrapper = PretrainTrainerWrapper(cfg, device="cpu")
|
|
mock_trainer = MagicMock()
|
|
mock_trainer.train = MagicMock()
|
|
mock_trainer.state.log_history = log_history if log_history is not None else []
|
|
mock_trainer.state.global_step = global_step
|
|
mock_trainer.save_model = MagicMock()
|
|
wrapper.trainer = mock_trainer
|
|
wrapper.tokenizer = MagicMock()
|
|
wrapper._output_dir = "./output"
|
|
return wrapper, mock_trainer
|
|
|
|
def test_train_returns_expected_keys(self):
|
|
"""train() result dict must contain all expected keys."""
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 3.0}, {"loss": 2.5}], global_step=10
|
|
)
|
|
result = wrapper.train()
|
|
assert "initial_loss" in result
|
|
assert "final_loss" in result
|
|
assert "duration" in result
|
|
assert "duration_secs" in result
|
|
assert "output_dir" in result
|
|
assert "total_steps" in result
|
|
|
|
def test_train_result_losses_from_log_history(self):
|
|
"""initial_loss and final_loss should come from trainer log_history."""
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 3.0}, {"loss": 2.5}, {"loss": 2.0}], global_step=30
|
|
)
|
|
result = wrapper.train()
|
|
assert result["initial_loss"] == pytest.approx(3.0)
|
|
assert result["final_loss"] == pytest.approx(2.0)
|
|
|
|
def test_train_result_empty_log_history_returns_zero_losses(self):
|
|
"""When log_history has no 'loss' entries, losses should be 0."""
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(log_history=[], global_step=0)
|
|
result = wrapper.train()
|
|
assert result["initial_loss"] == 0
|
|
assert result["final_loss"] == 0
|
|
|
|
def test_train_result_total_steps_from_trainer_state(self):
|
|
"""total_steps should match trainer.state.global_step."""
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 2.0}], global_step=42
|
|
)
|
|
result = wrapper.train()
|
|
assert result["total_steps"] == 42
|
|
|
|
def test_train_result_output_dir_matches(self):
|
|
"""output_dir in result should match wrapper._output_dir."""
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer()
|
|
result = wrapper.train()
|
|
assert result["output_dir"] == "./output"
|
|
|
|
def test_train_calls_save_model(self):
|
|
"""train() should call trainer.save_model with output_dir."""
|
|
wrapper, mock_trainer = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 2.0}], global_step=5
|
|
)
|
|
wrapper.train()
|
|
mock_trainer.save_model.assert_called_once_with("./output")
|
|
|
|
def test_train_calls_tokenizer_save_pretrained(self):
|
|
"""train() should call tokenizer.save_pretrained with output_dir."""
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 2.0}], global_step=5
|
|
)
|
|
wrapper.train()
|
|
wrapper.tokenizer.save_pretrained.assert_called_once_with("./output")
|
|
|
|
def test_train_passes_resume_checkpoint_to_trainer(self):
|
|
"""train() should forward resume_from_checkpoint to trainer.train()."""
|
|
wrapper, mock_trainer = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 2.0}], global_step=5
|
|
)
|
|
wrapper.train(resume_from_checkpoint="/ckpt/checkpoint-50")
|
|
mock_trainer.train.assert_called_once_with(
|
|
resume_from_checkpoint="/ckpt/checkpoint-50"
|
|
)
|
|
|
|
def test_train_result_duration_minutes_format(self):
|
|
"""Short durations (<1h) should produce 'Xm' format."""
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 2.0}], global_step=5
|
|
)
|
|
|
|
call_count = [0]
|
|
|
|
def fake_time():
|
|
call_count[0] += 1
|
|
return 0 if call_count[0] == 1 else 90
|
|
|
|
with mock_patch("soup_cli.trainer.pretrain.time.time", side_effect=fake_time):
|
|
result = wrapper.train()
|
|
|
|
assert result["duration"] == "1m"
|
|
|
|
def test_train_result_duration_hours_format(self):
|
|
"""Long durations (>=1h) should produce 'Xh Ym' format."""
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 2.0}], global_step=100
|
|
)
|
|
|
|
call_count = [0]
|
|
|
|
def fake_time():
|
|
call_count[0] += 1
|
|
return 0 if call_count[0] == 1 else 3720 # 1h 2m
|
|
|
|
with mock_patch("soup_cli.trainer.pretrain.time.time", side_effect=fake_time):
|
|
result = wrapper.train()
|
|
|
|
assert result["duration"] == "1h 2m"
|
|
|
|
|
|
# ─── CLI Init Template Tests ──────────────────────────────────────────────
|
|
|
|
|
|
class TestPretrainInitTemplate:
|
|
"""Test that soup init produces correct output for pretrain."""
|
|
|
|
def test_init_pretrain_template_creates_file(self, tmp_path):
|
|
"""soup init --template pretrain should write a file with pretrain task."""
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
|
|
runner = CliRunner()
|
|
output = tmp_path / "soup.yaml"
|
|
result = runner.invoke(
|
|
app, ["init", "--template", "pretrain", "--output", str(output)]
|
|
)
|
|
assert result.exit_code == 0
|
|
assert output.exists()
|
|
content = output.read_text()
|
|
assert "task: pretrain" in content
|
|
assert "format: plaintext" in content
|
|
|
|
def test_init_pretrain_template_produces_valid_config(self, tmp_path):
|
|
"""The file written by soup init --template pretrain should parse."""
|
|
from pathlib import Path
|
|
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
from soup_cli.config.loader import load_config
|
|
|
|
runner = CliRunner()
|
|
output = tmp_path / "soup.yaml"
|
|
runner.invoke(
|
|
app, ["init", "--template", "pretrain", "--output", str(output)]
|
|
)
|
|
cfg = load_config(Path(output))
|
|
assert cfg.task == "pretrain"
|
|
assert cfg.data.format == "plaintext"
|
|
|
|
def test_init_moe_template_creates_file(self, tmp_path):
|
|
"""soup init --template moe should write a file with moe_lora."""
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
|
|
runner = CliRunner()
|
|
output = tmp_path / "soup.yaml"
|
|
result = runner.invoke(
|
|
app, ["init", "--template", "moe", "--output", str(output)]
|
|
)
|
|
assert result.exit_code == 0
|
|
content = output.read_text()
|
|
assert "moe_lora: true" in content
|
|
|
|
|
|
# ─── Wizard Pretrain Path Tests ──────────────────────────────────────────
|
|
|
|
|
|
class TestPretrainWizardPath:
|
|
"""Test the interactive wizard auto-sets format for pretrain task."""
|
|
|
|
def test_wizard_pretrain_task_sets_plaintext_format(self):
|
|
"""When the wizard receives task=pretrain, data format should be 'plaintext'."""
|
|
from soup_cli.commands.init import _interactive_wizard
|
|
|
|
with mock_patch("soup_cli.commands.init.Prompt.ask", side_effect=[
|
|
"some-model",
|
|
"pretrain",
|
|
"./corpus.txt",
|
|
"1",
|
|
"yes",
|
|
]):
|
|
config_text = _interactive_wizard()
|
|
|
|
assert "task: pretrain" in config_text
|
|
assert "format: plaintext" in config_text
|
|
|
|
def test_wizard_pretrain_does_not_prompt_for_format(self):
|
|
"""The wizard should NOT ask for data format when task=pretrain."""
|
|
from soup_cli.commands.init import _interactive_wizard
|
|
|
|
prompt_calls = []
|
|
|
|
def record_prompt(question, **kwargs):
|
|
prompt_calls.append(question)
|
|
answers = {
|
|
"Base model": "some-model",
|
|
"Task": "pretrain",
|
|
"Training data path": "./corpus.txt",
|
|
"Epochs": "1",
|
|
"Use QLoRA (4-bit)?": "yes",
|
|
}
|
|
return answers.get(question, kwargs.get("default", ""))
|
|
|
|
with mock_patch("soup_cli.commands.init.Prompt.ask", side_effect=record_prompt):
|
|
config_text = _interactive_wizard()
|
|
|
|
# "Data format" prompt should not appear when task is pretrain
|
|
assert not any("format" in call.lower() for call in prompt_calls)
|
|
assert "format: plaintext" in config_text
|
|
|
|
|
|
# ─── Config Loader Round-trip Tests ──────────────────────────────────────
|
|
|
|
|
|
class TestPretrainConfigLoaderRoundTrip:
|
|
"""Test pretrain template YAML survives round-trip through load_config_from_string."""
|
|
|
|
def test_pretrain_template_round_trip(self):
|
|
"""TEMPLATES['pretrain'] should parse via load_config_from_string."""
|
|
from soup_cli.config.loader import load_config_from_string
|
|
|
|
cfg = load_config_from_string(TEMPLATES["pretrain"])
|
|
assert cfg.task == "pretrain"
|
|
assert cfg.data.format == "plaintext"
|
|
|
|
def test_moe_template_round_trip(self):
|
|
"""TEMPLATES['moe'] should parse via load_config_from_string."""
|
|
from soup_cli.config.loader import load_config_from_string
|
|
|
|
cfg = load_config_from_string(TEMPLATES["moe"])
|
|
assert cfg.training.moe_lora is True
|
|
assert cfg.training.moe_aux_loss_coeff == pytest.approx(0.01)
|
|
|
|
def test_pretrain_custom_yaml_round_trip(self):
|
|
"""Custom pretrain YAML string should round-trip correctly."""
|
|
from soup_cli.config.loader import load_config_from_string
|
|
|
|
yaml_str = """
|
|
base: custom-model/llama-7b
|
|
task: pretrain
|
|
|
|
data:
|
|
train: ./corpus.jsonl
|
|
format: plaintext
|
|
max_length: 4096
|
|
|
|
training:
|
|
epochs: 1
|
|
lr: 1e-5
|
|
quantization: none
|
|
|
|
output: ./pretrain_output
|
|
"""
|
|
cfg = load_config_from_string(yaml_str)
|
|
assert cfg.task == "pretrain"
|
|
assert cfg.data.format == "plaintext"
|
|
assert cfg.data.max_length == 4096
|
|
assert cfg.output == "./pretrain_output"
|
|
|
|
|
|
# ─── MoE Integration in _setup_transformers Tests ─────────────────────────
|
|
|
|
|
|
class TestPretrainMoEIntegration:
|
|
"""Test MoE-aware setup logic in PretrainTrainerWrapper._setup_transformers."""
|
|
|
|
def test_moe_lora_calls_get_moe_target_modules(self):
|
|
"""When moe_lora=True and model is MoE, get_moe_target_modules should be called."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
training={"moe_lora": True, "quantization": "none"},
|
|
)
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.get_nb_trainable_parameters.return_value = (1000, 10000)
|
|
|
|
with mock_patch("transformers.AutoModelForCausalLM.from_pretrained",
|
|
return_value=mock_model), \
|
|
mock_patch("transformers.AutoTokenizer.from_pretrained"), \
|
|
mock_patch("peft.get_peft_model", return_value=mock_model), \
|
|
mock_patch("peft.LoraConfig"), \
|
|
mock_patch("peft.prepare_model_for_kbit_training"), \
|
|
mock_patch(
|
|
"soup_cli.utils.moe.detect_moe_model", return_value=True
|
|
), \
|
|
mock_patch(
|
|
"soup_cli.utils.moe.get_moe_target_modules",
|
|
return_value=["q_proj", "v_proj", "gate_proj", "up_proj"],
|
|
) as mock_moe_targets:
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
wrapper = PretrainTrainerWrapper(cfg, device="cpu")
|
|
wrapper._setup_transformers(cfg, cfg.training)
|
|
|
|
mock_moe_targets.assert_called_once_with(mock_model)
|
|
|
|
def test_moe_aux_loss_coeff_sets_router_config(self):
|
|
"""When MoE detected and moe_aux_loss_coeff > 0, router config should be set."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
training={"moe_aux_loss_coeff": 0.05, "quantization": "none"},
|
|
)
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.get_nb_trainable_parameters.return_value = (1000, 10000)
|
|
mock_model.config.router_aux_loss_coef = 0.01
|
|
mock_model.config.output_router_logits = False
|
|
|
|
with mock_patch("transformers.AutoModelForCausalLM.from_pretrained",
|
|
return_value=mock_model), \
|
|
mock_patch("transformers.AutoTokenizer.from_pretrained"), \
|
|
mock_patch("peft.get_peft_model", return_value=mock_model), \
|
|
mock_patch("peft.LoraConfig"), \
|
|
mock_patch("peft.prepare_model_for_kbit_training"), \
|
|
mock_patch(
|
|
"soup_cli.utils.moe.detect_moe_model", return_value=True
|
|
), \
|
|
mock_patch("soup_cli.utils.moe.get_moe_target_modules",
|
|
return_value=None):
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
wrapper = PretrainTrainerWrapper(cfg, device="cpu")
|
|
wrapper._setup_transformers(cfg, cfg.training)
|
|
|
|
assert mock_model.config.router_aux_loss_coef == 0.05
|
|
assert mock_model.config.output_router_logits is True
|
|
|
|
def test_moe_lora_non_moe_model_skips_targets(self):
|
|
"""When moe_lora=True but model is not MoE, should fall back to auto."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="pretrain",
|
|
data={"train": "./data.jsonl"},
|
|
training={"moe_lora": True, "quantization": "none"},
|
|
)
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.get_nb_trainable_parameters.return_value = (1000, 10000)
|
|
|
|
with mock_patch("transformers.AutoModelForCausalLM.from_pretrained",
|
|
return_value=mock_model), \
|
|
mock_patch("transformers.AutoTokenizer.from_pretrained"), \
|
|
mock_patch("peft.get_peft_model", return_value=mock_model), \
|
|
mock_patch("peft.LoraConfig"), \
|
|
mock_patch("peft.prepare_model_for_kbit_training"), \
|
|
mock_patch(
|
|
"soup_cli.utils.moe.detect_moe_model", return_value=False
|
|
), \
|
|
mock_patch(
|
|
"soup_cli.utils.moe.get_moe_target_modules",
|
|
return_value=None,
|
|
) as mock_moe_targets:
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
wrapper = PretrainTrainerWrapper(cfg, device="cpu")
|
|
wrapper._setup_transformers(cfg, cfg.training)
|
|
|
|
# MoE targets should NOT be called since model is not MoE
|
|
mock_moe_targets.assert_not_called()
|
|
|
|
|
|
# ─── Plaintext Line-level Chunking Test ───────────────────────────────────
|
|
|
|
|
|
class TestPlaintextLineChunking:
|
|
"""Verify .txt files use line-level chunking (not paragraph-level)."""
|
|
|
|
def test_double_newlines_produce_separate_line_docs(self, tmp_path):
|
|
"""Double newlines are skipped — each non-empty line is a separate doc."""
|
|
from soup_cli.data.loader import load_raw_data
|
|
|
|
txt_file = tmp_path / "corpus.txt"
|
|
txt_file.write_text(
|
|
"Para one line one\n\nPara two line one\n", encoding="utf-8"
|
|
)
|
|
data = load_raw_data(txt_file)
|
|
assert len(data) == 2
|
|
assert data[0] == {"text": "Para one line one"}
|
|
assert data[1] == {"text": "Para two line one"}
|
|
|
|
def test_unicode_text_loading(self, tmp_path):
|
|
"""Unicode content should load correctly."""
|
|
from soup_cli.data.loader import load_raw_data
|
|
|
|
txt_file = tmp_path / "unicode.txt"
|
|
txt_file.write_text("日本語テスト\n中文测试\nعربي\n", encoding="utf-8")
|
|
data = load_raw_data(txt_file)
|
|
assert len(data) == 3
|
|
assert data[0] == {"text": "日本語テスト"}
|
|
assert data[1] == {"text": "中文测试"}
|
|
assert data[2] == {"text": "عربي"}
|