tests/test_training.py

"""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