File size: 4,036 Bytes
c63bc31 91c4b21 c63bc31 91c4b21 c63bc31 91c4b21 c63bc31 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Tests for the fine-tuning dataset pipeline and configuration."""
import json
from pathlib import Path
import pytest
from cortex_ai.training import (
DatasetError,
LoRAConfig,
TrainConfig,
dataset_stats,
iter_encoded,
load_jsonl,
normalize_record,
write_jsonl,
)
def test_normalize_chat_format():
ex = normalize_record(
{"messages": [{"role": "user", "content": "a"}, {"role": "assistant", "content": "b"}]}
)
assert ex.num_turns == 2
assert ex.messages[0]["role"] == "user"
def test_normalize_instruction_format():
ex = normalize_record({"instruction": "question", "response": "reponse", "system": "sys"})
assert [m["role"] for m in ex.messages] == ["system", "user", "assistant"]
assert ex.messages[2]["content"] == "reponse"
def test_normalize_rejects_unknown_shape():
with pytest.raises(DatasetError, match="expected"):
normalize_record({"foo": "bar"})
def test_normalize_rejects_bad_role():
with pytest.raises(DatasetError, match="invalid role"):
normalize_record({"messages": [{"role": "wizard", "content": "x"}]})
def test_normalize_rejects_empty_content():
with pytest.raises(DatasetError, match="empty content"):
normalize_record({"messages": [{"role": "user", "content": " "}]})
def test_normalize_rejects_dataset_without_assistant():
with pytest.raises(DatasetError, match="no assistant turn"):
normalize_record({"messages": [{"role": "user", "content": "x"}]})
def test_load_and_write_round_trip(tmp_path):
path = tmp_path / "d.jsonl"
path.write_text(
json.dumps({"instruction": "a", "response": "b"}) + "\n"
+ json.dumps({"messages": [{"role": "user", "content": "c"}, {"role": "assistant", "content": "d"}]}) + "\n",
encoding="utf-8",
)
examples = load_jsonl(path)
assert len(examples) == 2
out = tmp_path / "out.jsonl"
assert write_jsonl(examples, out) == 2
assert len(load_jsonl(out)) == 2
def test_load_jsonl_reports_bad_json(tmp_path):
path = tmp_path / "bad.jsonl"
path.write_text("{not json}\n", encoding="utf-8")
with pytest.raises(DatasetError, match="invalid JSON"):
load_jsonl(path)
def test_dataset_stats():
examples = load_jsonl(Path(__file__).parents[1] / "data" / "cortex_sft.jsonl")
stats = dataset_stats(examples)
assert stats["exemples"] == 10
assert stats["repartition_roles"]["assistant"] == 10
assert stats["caracteres_total"] > 0
def test_encoding_uses_the_real_encoder():
examples = load_jsonl(Path(__file__).parents[1] / "data" / "cortex_sft.jsonl")
prompts = list(iter_encoded(examples, thinking_mode="chat"))
assert len(prompts) == len(examples)
assert all(isinstance(p, str) and p for p in prompts)
def test_lora_config_validation():
with pytest.raises(ValueError, match="r must be"):
LoRAConfig(r=0)
with pytest.raises(ValueError, match="dropout"):
LoRAConfig(dropout=1.0)
with pytest.raises(ValueError, match="target_modules"):
LoRAConfig(target_modules=[])
assert LoRAConfig(r=8, alpha=16).scaling == 2.0
def test_train_config_validation():
with pytest.raises(ValueError, match="method"):
TrainConfig(method="magic")
with pytest.raises(ValueError, match="epochs"):
TrainConfig(epochs=0)
with pytest.raises(ValueError, match="thinking_mode"):
TrainConfig(thinking_mode="dream")
def test_effective_batch_size():
assert TrainConfig(batch_size=2, gradient_accumulation_steps=4).effective_batch_size == 8
def test_train_config_round_trip(tmp_path):
cfg = TrainConfig(method="lora", epochs=2, lora=LoRAConfig(r=8, alpha=8))
path = tmp_path / "train.json"
cfg.save(path)
restored = TrainConfig.load(path)
assert restored.epochs == 2
assert restored.lora.r == 8
assert restored.method == "lora"
def test_default_config_is_valid():
cfg = TrainConfig()
assert cfg.method == "lora"
assert cfg.effective_batch_size == 8
|