tests/test_finetune_lora_tooling.py

"""Tests for Modal LoRA fine-tuning scaffolding helpers."""

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from scripts import finetune_lora


class FakeTokenizer:
    pad_token = "<pad>"
    eos_token = "</s>"

    def apply_chat_template(
        self,
        messages: list[dict[str, str]],
        *,
        tokenize: bool,
        add_generation_prompt: bool,
    ) -> str:
        text = "".join(
            f"<{message['role']}>{message['content']}</{message['role']}>"
            for message in messages
        )
        if add_generation_prompt:
            text += "<assistant>"
        return text

    def __call__(
        self,
        text: str,
        *,
        truncation: bool,
        max_length: int,
        padding: bool,
        add_special_tokens: bool = False,
    ) -> dict[str, list[int]]:
        del padding, add_special_tokens
        ids = [ord(character) % 251 + 1 for character in text]
        if truncation:
            ids = ids[:max_length]
        return {"input_ids": ids, "attention_mask": [1] * len(ids)}


def _valid_record() -> dict[str, object]:
    return {
        "id": "sft-preview-0001",
        "messages": [
            {"role": "system", "content": "You are Objectverse Diary."},
            {"role": "user", "content": "Create a persona."},
            {"role": "assistant", "content": "{\"persona\": {}, \"diary\": {}}"},
        ],
    }


class FinetuneLoraToolingTest(unittest.TestCase):
    def test_load_sft_records_rejects_missing_messages(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            path = Path(tmp_dir) / "bad.jsonl"
            path.write_text(json.dumps({"id": "bad"}) + "\n", encoding="utf-8")

            with self.assertRaises(ValueError):
                finetune_lora.load_sft_records(path)

    def test_load_sft_records_rejects_malformed_messages(self) -> None:
        bad_record = {"id": "bad", "messages": [{"role": "user"}]}
        with tempfile.TemporaryDirectory() as tmp_dir:
            path = Path(tmp_dir) / "bad.jsonl"
            path.write_text(json.dumps(bad_record) + "\n", encoding="utf-8")

            with self.assertRaises(ValueError):
                finetune_lora.load_sft_records(path)

    def test_record_to_training_text_is_non_empty(self) -> None:
        text = finetune_lora.record_to_training_text(_valid_record())

        self.assertIn("system:", text)
        self.assertIn("user:", text)
        self.assertIn("assistant:", text)
        self.assertIn("Objectverse Diary", text)

    def test_default_training_config_uses_safe_qwen_lora_defaults(self) -> None:
        config = finetune_lora.TrainingConfig()

        self.assertEqual(config.base_model, "Qwen/Qwen2.5-1.5B-Instruct")
        self.assertEqual(config.lora_r, 16)
        self.assertEqual(config.lora_alpha, 32)
        self.assertEqual(config.lora_dropout, 0.05)
        self.assertEqual(config.max_steps, 80)
        self.assertEqual(config.num_train_epochs, 3.0)
        self.assertEqual(config.per_device_train_batch_size, 1)
        self.assertEqual(config.gradient_accumulation_steps, 4)
        self.assertEqual(config.eval_ratio, 0.1)
        self.assertTrue(config.assistant_only_loss)
        self.assertIn("q_proj", config.target_modules)
        self.assertIn("down_proj", config.target_modules)

    def test_training_config_serializes_v2_experiment_settings(self) -> None:
        config = finetune_lora.TrainingConfig(
            max_steps=0,
            num_train_epochs=4.0,
            per_device_train_batch_size=2,
            gradient_accumulation_steps=8,
            eval_ratio=0.2,
            eval_steps=25,
            lora_r=32,
            lora_alpha=64,
            assistant_only_loss=False,
        )

        payload = config.as_remote_dict()

        self.assertEqual(payload["num_train_epochs"], 4.0)
        self.assertEqual(payload["per_device_train_batch_size"], 2)
        self.assertEqual(payload["gradient_accumulation_steps"], 8)
        self.assertEqual(payload["eval_ratio"], 0.2)
        self.assertEqual(payload["eval_steps"], 25)
        self.assertEqual(payload["lora_r"], 32)
        self.assertEqual(payload["lora_alpha"], 64)
        self.assertFalse(payload["assistant_only_loss"])

    def test_dry_run_does_not_call_remote_runner(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            path = Path(tmp_dir) / "records.jsonl"
            path.write_text(json.dumps(_valid_record()) + "\n", encoding="utf-8")

            def fail_remote_call(
                records: list[dict[str, object]],
                config: finetune_lora.TrainingConfig,
            ) -> dict[str, object]:
                raise AssertionError("dry-run should not call remote training")

            summary = finetune_lora.run_training_entrypoint(
                dataset=path,
                config=finetune_lora.TrainingConfig(),
                dry_run=True,
                allow_remote=False,
                remote_runner=fail_remote_call,
            )

        self.assertEqual(summary["mode"], "dry-run")
        self.assertEqual(summary["record_count"], 1)
        self.assertEqual(summary["base_model"], "Qwen/Qwen2.5-1.5B-Instruct")
        self.assertEqual(summary["train_record_count"], 1)
        self.assertEqual(summary["eval_record_count"], 0)

    def test_dry_run_reports_eval_split_for_larger_datasets(self) -> None:
        records = [_valid_record() for _ in range(20)]
        summary = finetune_lora._dry_run_summary(
            Path("records.jsonl"),
            records,
            finetune_lora.TrainingConfig(eval_ratio=0.2),
        )

        self.assertEqual(summary["train_record_count"], 16)
        self.assertEqual(summary["eval_record_count"], 4)

    def test_assistant_only_tokenization_masks_prompt_labels(self) -> None:
        tokenized = finetune_lora._tokenize_training_example(
            _valid_record(),
            FakeTokenizer(),
            max_length=512,
            assistant_only_loss=True,
        )

        labels = tokenized["labels"]

        self.assertIn(-100, labels)
        self.assertTrue(any(label != -100 for label in labels))
        first_unmasked = next(index for index, label in enumerate(labels) if label != -100)
        self.assertGreater(first_unmasked, 0)

    def test_full_loss_tokenization_keeps_all_labels(self) -> None:
        tokenized = finetune_lora._tokenize_training_example(
            _valid_record(),
            FakeTokenizer(),
            max_length=512,
            assistant_only_loss=False,
        )

        self.assertNotIn(-100, tokenized["labels"])


if __name__ == "__main__":
    unittest.main()