apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-26 17:11:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

icml rebuttal

Browse source
This commit is contained in:
51616 2026-06-15 04:31:47 +00:00
parent 22267c7666
commit 696b45c51b
44 changed files with 5300 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

91
tests/test_chunk_overlap.py Normal file
View file

@ -0,0 +1,91 @@
import inspect
import unittest
from ctx_to_lora.data.definitions import CTX_AFFIXES
from ctx_to_lora.data.processing import (
construct_and_tokenize_ctx_qa,
split_too_long_ctx,
)
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
class ChunkOverlapTest(unittest.TestCase):
def test_construct_and_tokenize_ctx_qa_accepts_overlap_kwarg(self):
signature = inspect.signature(construct_and_tokenize_ctx_qa)
self.assertIn("ctx_chunk_overlap", signature.parameters)
def _strip_affixes(self, chunks):
prefix = CTX_AFFIXES[MODEL_NAME]["prefix"]
suffix = CTX_AFFIXES[MODEL_NAME]["suffix"]
out = []
for idx, chunk in enumerate(chunks):
if idx > 0:
self.assertEqual(chunk[: len(prefix)], prefix)
chunk = chunk[len(prefix) :]
if idx < len(chunks) - 1:
self.assertEqual(chunk[-len(suffix) :], suffix)
chunk = chunk[: -len(suffix)]
out.append(chunk)
return out
def test_eval_overlap_adds_shared_boundary_tokens(self):
ctx_ids = list(range(10))
out = split_too_long_ctx(
sample={"ctx_ids": ctx_ids},
model_name_or_path=MODEL_NAME,
num_chunk_probs=None,
max_chunk_len=4,
min_chunk_len=-1,
max_num_split=None,
is_train=False,
chunk_overlap=2,
)
stripped_chunks = self._strip_affixes(out["ctx_ids"])
self.assertEqual(out["n_ctx_chunks"], 4)
self.assertEqual(
stripped_chunks,
[
[0, 1, 2, 3],
[2, 3, 4, 5],
[4, 5, 6, 7],
[6, 7, 8, 9],
],
)
def test_zero_overlap_preserves_existing_balanced_split(self):
ctx_ids = list(range(10))
out = split_too_long_ctx(
sample={"ctx_ids": ctx_ids},
model_name_or_path=MODEL_NAME,
num_chunk_probs=None,
max_chunk_len=4,
min_chunk_len=-1,
max_num_split=None,
is_train=False,
chunk_overlap=0,
)
stripped_chunks = self._strip_affixes(out["ctx_ids"])
self.assertEqual(out["n_ctx_chunks"], 3)
self.assertEqual(stripped_chunks, [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]])
def test_overlap_is_rejected_for_train(self):
with self.assertRaisesRegex(ValueError, "eval splits"):
split_too_long_ctx(
sample={"ctx_ids": list(range(8))},
model_name_or_path=MODEL_NAME,
num_chunk_probs=None,
max_chunk_len=4,
min_chunk_len=-1,
max_num_split=None,
is_train=True,
chunk_overlap=1,
)
if __name__ == "__main__":
unittest.main()

60
tests/test_clipper_eval.py Normal file
View file

@ -0,0 +1,60 @@
from ctx_to_lora.data.preprocessing_fn import parse_clipper_user_message
from ctx_to_lora.eval_utils import compute_clipper_metrics, extract_clipper_answer
def test_parse_clipper_user_message_splits_context_and_prompt():
user_message = (
"You are provided with a context and a statement.\n\n"
"<context>Book body</context>\n\n"
"<statement>A claim.</statement>\n\n"
"<question>Is it true or false?</question>"
)
context, prompt = parse_clipper_user_message(user_message)
assert context == "Book body"
assert "<statement>A claim.</statement>" in prompt
assert "<context>" not in prompt
def test_extract_clipper_answer_prefers_answer_tag():
text = "<explanation>Reasoning mentions false.</explanation><answer>TRUE</answer>"
assert extract_clipper_answer(text) == "true"
def test_compute_clipper_metrics_reports_pair_accuracy():
decoded_txts = [
{
"generated": "<answer>TRUE</answer>",
"label": "<answer>TRUE</answer>",
"clipper_status": "true",
"clipper_pair_id": "pair-a",
},
{
"generated": "<answer>false</answer>",
"label": "<answer>FALSE</answer>",
"clipper_status": "false",
"clipper_pair_id": "pair-a",
},
{
"generated": "<answer>true</answer>",
"label": "<answer>TRUE</answer>",
"clipper_status": "true",
"clipper_pair_id": "pair-b",
},
{
"generated": "<answer>true</answer>",
"label": "<answer>FALSE</answer>",
"clipper_status": "false",
"clipper_pair_id": "pair-b",
},
]
metrics, per_sample, counts = compute_clipper_metrics(decoded_txts)
assert metrics["clipper_accuracy"] == 0.75
assert metrics["clipper_true_accuracy"] == 1.0
assert metrics["clipper_false_accuracy"] == 0.5
assert metrics["clipper_pair_accuracy"] == 0.5
assert per_sample["clipper_pair_accuracy"] == [1.0, 1.0, 0.0, 0.0]
assert counts["clipper_pair_accuracy"] == 2

249
tests/test_dense_lora_magnitude.py Normal file
View file

@ -0,0 +1,249 @@
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.eval_utils import eval_generation
from ctx_to_lora.lora_stats import DenseLoraMagnitudeAggregator
def _manual_dense_stats(combined_loras, scaling):
overall = []
by_module = {}
for module_name, module_lora in combined_loras.items():
module_values = []
A = module_lora["A"]
B = module_lora["B"]
for batch_idx in range(A.shape[0]):
for layer_idx in range(A.shape[1]):
dense = (
torch.einsum(
"ro,ri->oi",
B[batch_idx, layer_idx],
A[batch_idx, layer_idx],
)
* scaling
)
module_values.append(dense.abs().reshape(-1))
by_module[module_name] = torch.cat(module_values)
overall.append(by_module[module_name])
return torch.cat(overall), by_module
class DenseLoraMagnitudeTests(unittest.TestCase):
@unittest.skipUnless(
torch.cuda.is_available(),
"CUDA is required for dense LoRA magnitude computation.",
)
def test_dense_lora_magnitude_aggregator_matches_manual_batched_dense_update(self):
scaling = 1.5
combined_loras = {
"down_proj": {
"A": torch.tensor(
[
[
[[1.0, -2.0], [0.5, 3.0]],
[[-1.0, 2.0], [4.0, -0.5]],
],
[
[[2.0, 1.0], [1.5, -1.0]],
[[0.25, -3.0], [2.0, 0.75]],
],
]
),
"B": torch.tensor(
[
[
[[2.0, -1.0, 0.5], [1.0, 0.0, -2.0]],
[[0.5, 1.5, -1.0], [2.0, -0.5, 1.0]],
],
[
[[-1.0, 2.0, 0.25], [0.5, -1.5, 3.0]],
[[1.0, 0.0, 2.0], [-2.0, 1.0, -0.5]],
],
]
),
},
"up_proj": {
"A": torch.tensor(
[
[
[[0.5, 1.0], [2.0, -1.0]],
[[1.0, -0.5], [0.25, 3.0]],
],
[
[[-2.0, 0.5], [1.0, 1.5]],
[[0.75, -1.25], [2.5, 0.5]],
],
]
),
"B": torch.tensor(
[
[
[[1.0, -2.0], [0.5, 1.5]],
[[-1.0, 2.0], [3.0, 0.25]],
],
[
[[2.0, 1.0], [-0.5, 0.5]],
[[1.5, -1.0], [0.0, 2.0]],
],
]
),
},
}
aggregator = DenseLoraMagnitudeAggregator(scaling=scaling, out_chunk_size=2)
aggregator.update(combined_loras)
overall_manual, by_module_manual = _manual_dense_stats(combined_loras, scaling)
stats = aggregator.to_dict()
self.assertEqual(stats["num_examples"], 2)
self.assertEqual(stats["num_batches"], 1)
self.assertEqual(stats["num_layer_matrices"], 4)
self.assertAlmostEqual(
stats["overall"]["mean"], overall_manual.mean().item(), places=6
)
self.assertAlmostEqual(
stats["overall"]["min"], overall_manual.min().item(), places=6
)
self.assertAlmostEqual(
stats["overall"]["max"], overall_manual.max().item(), places=6
)
self.assertAlmostEqual(
stats["overall"]["std"],
overall_manual.std(unbiased=False).item(),
places=6,
)
for module_name, module_values in by_module_manual.items():
module_stats = stats["by_module"][module_name]
self.assertAlmostEqual(
module_stats["mean"], module_values.mean().item(), places=6
)
self.assertAlmostEqual(
module_stats["min"], module_values.min().item(), places=6
)
self.assertAlmostEqual(
module_stats["max"], module_values.max().item(), places=6
)
self.assertAlmostEqual(
module_stats["std"], module_values.std(unbiased=False).item(), places=6
)
def test_eval_generation_saves_dense_lora_magnitude_json_and_metrics(self):
dense_stats = {
"overall": {"count": 8, "mean": 1.25, "min": 0.0, "max": 4.0, "std": 1.0},
"by_module": {
"down_proj": {
"count": 8,
"mean": 1.25,
"min": 0.0,
"max": 4.0,
"std": 1.0,
}
},
"num_examples": 1,
"num_batches": 1,
"num_layer_matrices": 2,
"out_chunk_size": 8,
"scaling": 16.0,
}
class FakeModel:
def __init__(self, stats):
self.log_dense_lora_magnitude = True
self._stats = stats
self.reset_calls = 0
def reset_dense_lora_magnitude_stats(self):
self.reset_calls += 1
def get_dense_lora_magnitude_stats(self):
return self._stats
class FakeTrainer:
def __init__(self, output_dir, stats):
self.args = SimpleNamespace(output_dir=output_dir)
self.model = FakeModel(stats)
self.logged = []
self.saved = []
def predict(self, *args, **kwargs):
return SimpleNamespace(predictions=[[0, 1]], metrics={})
def log_metrics(self, split, metrics):
self.logged.append((split, metrics))
def save_metrics(self, split, metrics):
self.saved.append((split, metrics))
with tempfile.TemporaryDirectory() as tmpdir:
trainer = FakeTrainer(tmpdir, dense_stats)
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
):
result = 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},
)
json_path = Path(tmpdir) / "test_dummy_dense_lora_magnitude.json"
self.assertTrue(json_path.exists())
self.assertEqual(json.loads(json_path.read_text()), dense_stats)
self.assertEqual(trainer.model.reset_calls, 1)
self.assertIn("test_dummy_dense_lora_magnitude_mean", result["test_dummy"])
self.assertEqual(
result["test_dummy"]["test_dummy_dense_lora_magnitude_mean"], 1.25
)
self.assertIn(
"test_dummy_dense_lora_magnitude_down_proj_max", result["test_dummy"]
)
def test_run_eval_cli_forwards_log_dense_lora_magnitude_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),
"--model_name_or_path",
"google/gemma-2-2b-it",
"--datasets",
"dummy_dataset",
"--log_dense_lora_magnitude",
]
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["log_dense_lora_magnitude"])
self.assertTrue(captured["generative"])
if __name__ == "__main__":
unittest.main()

110
tests/test_oolong_eval.py Normal file
View file

@ -0,0 +1,110 @@
from collections import Counter
from datasets import Dataset
from ctx_to_lora.data.preprocessing_fn import get_preprocessing_fn
from ctx_to_lora.data.processing import (
OOLONG_MAX_CONTEXT_LEN,
filter_dataset_by_max_value,
stratified_subsample_dataset,
)
from ctx_to_lora.eval_utils import compute_oolong_metrics, extract_oolong_answer
def test_oolong_preprocessing_maps_core_fields_and_metadata():
preprocess = get_preprocessing_fn("oolong-synth", is_eval=True)
sample = {
"id": 7,
"dataset": "spam",
"context_window_text": "context body",
"question": "Give your final answer in the form 'Label: answer'.",
"task_group": "classification",
"task": "most_freq",
"answer": "['ham']",
"answer_type": "ANSWER_TYPE.LABEL",
"input_subset": "all",
"num_labels": 2,
"context_window_id": 101,
}
processed = preprocess(sample)
assert processed["context"] == "context body"
assert processed["prompts"] == [sample["question"]]
assert processed["responses"] == ["['ham']"]
assert processed["oolong_task"] == "most_freq"
assert processed["oolong_answer_type"] == "ANSWER_TYPE.LABEL"
assert processed["oolong_context_window_id"] == 101
def test_extract_oolong_answer_prefers_suffix_after_colon():
answer, confidence = extract_oolong_answer("Reasoning...\nLabel: [ham]")
assert answer == "ham"
assert confidence == "vhigh"
def test_compute_oolong_metrics_handles_exact_numeric_and_date_scoring():
decoded_txts = [
{
"generated": "Label: ham",
"label": "['ham']",
"oolong_answer_type": "ANSWER_TYPE.LABEL",
},
{
"generated": "Answer: 8",
"label": "['10']",
"oolong_answer_type": "ANSWER_TYPE.NUMERIC",
},
{
"generated": "Date: 2025-02-07",
"label": "[datetime.date(2025, 2, 7)]",
"oolong_answer_type": "ANSWER_TYPE.DATE",
},
]
metrics, per_sample, counts = compute_oolong_metrics(decoded_txts)
assert metrics["oolong_score"] == (1.0 + (0.75**2) + 1.0) / 3
assert metrics["oolong_parse_rate"] == 1.0
assert per_sample["oolong_score"] == [1.0, 0.75**2, 1.0]
assert counts["oolong_score"] == 3
def test_stratified_subsample_dataset_balances_context_len_groups():
ds = Dataset.from_dict(
{
"context_len": [1] * 800 + [2] * 50 + [3] * 800,
"row_id": list(range(1650)),
}
)
sampled = stratified_subsample_dataset(
ds=ds,
group_column="context_len",
max_samples=1000,
seed=42,
)
counts = Counter(sampled["context_len"])
assert len(sampled) == 1000
assert counts[2] == 50
assert counts[1] == 475
assert counts[3] == 475
def test_filter_dataset_by_max_value_drops_over_64k_context_len_rows():
ds = Dataset.from_dict(
{
"context_len": [1024, OOLONG_MAX_CONTEXT_LEN, OOLONG_MAX_CONTEXT_LEN + 1],
"row_id": [0, 1, 2],
}
)
filtered = filter_dataset_by_max_value(
ds=ds,
column="context_len",
max_value=OOLONG_MAX_CONTEXT_LEN,
)
assert filtered["context_len"] == [1024, OOLONG_MAX_CONTEXT_LEN]

410
tests/test_rag.py Normal file
View file

@ -0,0 +1,410 @@
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()