doc-to-lora/tests/test_rag.py

import json
import runpy
import sys
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import torch

from ctx_to_lora.data.definitions import CTX_AFFIXES
from ctx_to_lora.eval_utils import eval_generation, run_eval
from ctx_to_lora.modeling.rag import DocToLoraRAGModel, RAGModel, RetrievedChunk

MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
PREFIX = CTX_AFFIXES[MODEL_NAME]["prefix"]
SUFFIX = CTX_AFFIXES[MODEL_NAME]["suffix"]


class FakeTokenizer:
    def __init__(self, name_or_path=MODEL_NAME):
        self.name_or_path = name_or_path
        self.pad_token_id = 0
        self.eos_token_id = 1
        self.chat_template = "fake"
        self._token_to_id = {}
        self._id_to_token = {}
        self._next_id = 1000

    def _encode_words(self, text: str) -> list[int]:
        tokens = []
        for word in text.split():
            if word not in self._token_to_id:
                token_id = self._next_id
                self._next_id += 1
                self._token_to_id[word] = token_id
                self._id_to_token[token_id] = word
            tokens.append(self._token_to_id[word])
        return tokens

    def __call__(
        self,
        text,
        add_special_tokens=False,
        return_attention_mask=False,
        return_tensors=None,
        **kwargs,
    ):
        del add_special_tokens, kwargs
        input_ids = self._encode_words(text)
        out = {"input_ids": input_ids}
        if return_attention_mask:
            out["attention_mask"] = [1] * len(input_ids)
        if return_tensors == "pt":
            out = {k: torch.tensor([v], dtype=torch.long) for k, v in out.items()}
        return out

    def decode(self, ids, skip_special_tokens=True):
        del skip_special_tokens
        if isinstance(ids, torch.Tensor):
            ids = ids.tolist()
        words = []
        for token_id in ids:
            if token_id in self._id_to_token:
                words.append(self._id_to_token[token_id])
        return " ".join(words)

    def apply_chat_template(
        self,
        conversations,
        tokenize=True,
        add_generation_prompt=True,
        return_attention_mask=True,
        return_dict=True,
        return_tensors=None,
        add_special_tokens=False,
        **kwargs,
    ):
        del tokenize, add_special_tokens, kwargs
        messages = (
            conversations[0]
            if conversations and isinstance(conversations[0], list)
            else conversations
        )
        user_content = [m["content"] for m in messages if m["role"] == "user"][-1]
        token_ids = PREFIX + self._encode_words(user_content)
        if add_generation_prompt:
            token_ids += SUFFIX
        out = {"input_ids": [token_ids]}
        if return_attention_mask:
            out["attention_mask"] = [[1] * len(token_ids)]
        if return_tensors == "pt":
            out = {k: torch.tensor(v, dtype=torch.long) for k, v in out.items()}
        if return_dict:
            return out
        return out["input_ids"]


class FakeBaseModel(torch.nn.Module):
    def __init__(self, name_or_path=MODEL_NAME):
        super().__init__()
        self.dummy = torch.nn.Parameter(torch.zeros(1))
        self.name_or_path = name_or_path
        self.config = SimpleNamespace(
            name_or_path=name_or_path,
            _name_or_path=name_or_path,
            max_position_embeddings=4096,
        )
        self.generation_config = SimpleNamespace(pad_token_id=0)
        self.last_generate_kwargs = None

    def generate(self, *args, **kwargs):
        del args
        self.last_generate_kwargs = kwargs
        return kwargs["input_ids"]


class FakeDocToLoraModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.base_model = FakeBaseModel()
        self.generation_config = self.base_model.generation_config
        self.last_generate_kwargs = None

    def generate(self, *args, **kwargs):
        del args
        self.last_generate_kwargs = dict(kwargs)
        return kwargs["input_ids"]


def build_chat_tensor(tokenizer: FakeTokenizer, text: str) -> dict[str, torch.Tensor]:
    token_ids = PREFIX + tokenizer._encode_words(text) + SUFFIX
    return {
        "input_ids": torch.tensor([token_ids], dtype=torch.long),
        "attention_mask": torch.ones((1, len(token_ids)), dtype=torch.long),
    }


class RAGModelTests(unittest.TestCase):
    def setUp(self):
        self.tokenizer = FakeTokenizer()
        self.base_model = FakeBaseModel()
        self.rag = RAGModel(
            model=self.base_model,
            tokenizer=self.tokenizer,
            chunk_size=4,
            chunk_overlap=2,
            top_k=2,
            max_retrieved_tokens=8,
        )

    def test_chunk_context_respects_overlap(self):
        chunks = self.rag._chunk_context("a b c d e f g h i j")
        self.assertEqual(
            [(chunk.start, chunk.end) for chunk in chunks],
            [(0, 4), (2, 6), (4, 8), (6, 10)],
        )
        self.assertEqual(chunks[1].text, "c d e f")
        self.assertEqual(chunks[-1].text, "g h i j")

    def test_bm25_prefers_relevant_chunk(self):
        chunks = [
            RetrievedChunk(
                "apple orange", self.tokenizer._encode_words("apple orange"), 0, 2, 0.0
            ),
            RetrievedChunk(
                "banana pear", self.tokenizer._encode_words("banana pear"), 2, 4, 0.0
            ),
        ]
        ranked = self.rag._score_chunks_bm25("apple", chunks)
        self.assertEqual(ranked[0].text, "apple orange")
        self.assertGreater(ranked[0].score, ranked[1].score)

    def test_generate_rebuilds_prompt_from_selected_chunks(self):
        rag = RAGModel(
            model=self.base_model,
            tokenizer=self.tokenizer,
            chunk_size=3,
            chunk_overlap=0,
            top_k=1,
            max_retrieved_tokens=3,
        )
        question_inputs = build_chat_tensor(self.tokenizer, "answer")
        context_inputs = build_chat_tensor(
            self.tokenizer,
            "cat dog fish bird apple answer kiwi mango noise noise",
        )

        rag.generate(
            input_ids=question_inputs["input_ids"],
            attention_mask=question_inputs["attention_mask"],
            ctx_ids=context_inputs["input_ids"],
            ctx_attn_mask=context_inputs["attention_mask"],
            max_new_tokens=16,
        )

        generated_input_ids = self.base_model.last_generate_kwargs["input_ids"][0]
        generated_attention_mask = self.base_model.last_generate_kwargs[
            "attention_mask"
        ][0]
        rebuilt_prompt, _ = rag._decode_question(
            generated_input_ids,
            generated_attention_mask,
        )

        self.assertIn("bird apple answer", rebuilt_prompt)
        self.assertIn("Question:", rebuilt_prompt)
        self.assertNotIn("cat dog fish", rebuilt_prompt)
        self.assertEqual(len(rag.get_retrieval_history()), 1)
        self.assertEqual(
            rag.get_retrieval_history()[0]["rag_selected_chunks"][0]["text"],
            "bird apple answer",
        )

    def test_hybrid_generate_preserves_ctx_and_rewrites_prompt_input(self):
        doc_model = FakeDocToLoraModel()
        hybrid = DocToLoraRAGModel(
            model=doc_model,
            tokenizer=self.tokenizer,
            chunk_size=3,
            chunk_overlap=0,
            top_k=1,
            max_retrieved_tokens=3,
        )
        question_inputs = build_chat_tensor(self.tokenizer, "answer")
        context_inputs = build_chat_tensor(
            self.tokenizer,
            "cat dog fish bird apple answer kiwi mango noise noise",
        )

        hybrid.generate(
            input_ids=question_inputs["input_ids"],
            attention_mask=question_inputs["attention_mask"],
            ctx_ids=context_inputs["input_ids"],
            ctx_attn_mask=context_inputs["attention_mask"],
            n_ctx_chunks=torch.tensor([1], dtype=torch.int32),
            max_new_tokens=8,
        )

        forwarded = doc_model.last_generate_kwargs
        self.assertTrue(torch.equal(forwarded["ctx_ids"], context_inputs["input_ids"]))
        self.assertTrue(
            torch.equal(forwarded["ctx_attn_mask"], context_inputs["attention_mask"])
        )
        self.assertTrue(
            torch.equal(forwarded["n_ctx_chunks"], torch.tensor([1], dtype=torch.int32))
        )
        rebuilt_prompt, _ = hybrid._decode_question(
            forwarded["input_ids"][0],
            forwarded["attention_mask"][0],
        )
        self.assertIn("bird apple answer", rebuilt_prompt)
        self.assertEqual(
            hybrid.get_retrieval_history()[0]["rag_selected_chunks"][0]["text"],
            "bird apple answer",
        )

    def test_run_eval_rejects_eval_chunking_with_rag(self):
        with self.assertRaisesRegex(ValueError, "retrieval chunking"):
            run_eval(
                model_name_or_path="google/gemma-2-2b-it",
                datasets=["dummy_dataset"],
                use_rag=True,
                max_ctx_chunk_len=128,
                eval_batch_size=1,
            )

    def test_run_eval_rejects_hybrid_without_checkpoint(self):
        with self.assertRaisesRegex(ValueError, "only available"):
            run_eval(
                model_name_or_path="google/gemma-2-2b-it",
                datasets=["dummy_dataset"],
                use_hybrid_rag=True,
                eval_batch_size=1,
            )

    def test_eval_generation_saves_retrieval_metadata(self):
        class FakeModel:
            def __init__(self):
                self._history = [
                    {
                        "rag_question": "who?",
                        "rag_selected_chunks": [
                            {"text": "retrieved text", "score": 1.0}
                        ],
                    }
                ]

            def reset_retrieval_history(self):
                self._history = list(self._history)

            def get_retrieval_history(self):
                return self._history

        class FakeTrainer:
            def __init__(self, output_dir):
                self.args = SimpleNamespace(output_dir=output_dir)
                self.model = FakeModel()

            def predict(self, *args, **kwargs):
                del args, kwargs
                return SimpleNamespace(predictions=[[0, 1]], metrics={})

            def log_metrics(self, split, metrics):
                del split, metrics

            def save_metrics(self, split, metrics):
                del split, metrics

        with tempfile.TemporaryDirectory() as tmpdir:
            trainer = FakeTrainer(tmpdir)
            decoded = [{"generated": "answer", "label": "answer", "ctx_ids_len": 4}]
            dataset = [{"input_ids": [1], "labels": [-100, 1]}]

            with patch(
                "ctx_to_lora.eval_utils.decode_test_result", return_value=decoded
            ):
                eval_generation(
                    trainer,
                    tokenizer=None,
                    ctx_tokenizer=None,
                    datasets={"dummy": dataset},
                    original_datasets={},
                    answers={},
                    split="test",
                    remove_context=False,
                    gen_kwargs={"max_new_tokens": 4, "do_sample": False},
                )

            jsonl_path = Path(tmpdir) / "test_dummy_generated_text.jsonl"
            lines = jsonl_path.read_text().splitlines()
            payload = json.loads(lines[0])
            self.assertEqual(payload["rag_question"], "who?")
            self.assertEqual(
                payload["rag_selected_chunks"][0]["text"], "retrieved text"
            )

    def test_run_eval_cli_forwards_rag_flags(self):
        script_path = Path(__file__).resolve().parents[1] / "run_eval.py"
        captured = {}

        def fake_run_eval(**kwargs):
            captured.update(kwargs)

        argv = [
            str(script_path),
            "--model_name_or_path",
            "google/gemma-2-2b-it",
            "--datasets",
            "dummy_dataset",
            "--use_rag",
            "--rag_chunk_size",
            "128",
            "--rag_chunk_overlap",
            "16",
            "--rag_top_k",
            "3",
            "--rag_max_retrieved_tokens",
            "512",
        ]

        with patch("ctx_to_lora.eval_utils.run_eval", side_effect=fake_run_eval):
            old_argv = sys.argv
            try:
                sys.argv = argv
                runpy.run_path(str(script_path), run_name="__main__")
            finally:
                sys.argv = old_argv

        self.assertTrue(captured["use_rag"])
        self.assertEqual(captured["rag_chunk_size"], 128)
        self.assertEqual(captured["rag_chunk_overlap"], 16)
        self.assertEqual(captured["rag_top_k"], 3)
        self.assertEqual(captured["rag_max_retrieved_tokens"], 512)
        self.assertTrue(captured["generative"])

    def test_run_eval_cli_forwards_hybrid_rag_flag(self):
        script_path = Path(__file__).resolve().parents[1] / "run_eval.py"
        captured = {}

        def fake_run_eval(**kwargs):
            captured.update(kwargs)

        argv = [
            str(script_path),
            "--checkpoint_path",
            "train_outputs/runs/demo/checkpoint-1/pytorch_model.bin",
            "--datasets",
            "dummy_dataset",
            "--use_hybrid_rag",
            "--rag_chunk_size",
            "128",
        ]

        with patch("ctx_to_lora.eval_utils.run_eval", side_effect=fake_run_eval):
            old_argv = sys.argv
            try:
                sys.argv = argv
                runpy.run_path(str(script_path), run_name="__main__")
            finally:
                sys.argv = old_argv

        self.assertTrue(captured["use_hybrid_rag"])
        self.assertEqual(captured["rag_chunk_size"], 128)
        self.assertTrue(captured["generative"])


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