mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
410 lines
14 KiB
Python
410 lines
14 KiB
Python
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()
|