File size: 4,946 Bytes
20c251e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | from __future__ import annotations
from pathlib import Path
from dovla_cil.data.datasets import CILDataset
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.retrieval.embeddings import cosine_similarity, embed_observation_language
from dovla_cil.retrieval.eval import RetrievalEvalQuery, evaluate_retrieval_baselines
from dovla_cil.retrieval.index import CILRetrievalIndex
from dovla_cil.retrieval.prompting import build_retrieval_prompt
from dovla_cil.retrieval.retriever import CriticGatedRetriever, RetrievalConditionedPolicyWrapper
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.transfercritic.schema import TransferContext
def _make_dataset(tmp_path: Path) -> CILDataset:
generate_cil_dataset(
backend="toy",
tasks=built_in_toy_tasks()[:2],
out_dir=tmp_path,
num_states_per_task=1,
k=4,
seed=21,
shard_size=8,
inline_observations=True,
)
return CILDataset(tmp_path)
def test_embedding_index_over_tiny_cil_dataset(tmp_path: Path) -> None:
dataset = _make_dataset(tmp_path)
index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
record = dataset[0]
query = embed_observation_language(record.observation_inline, record.instruction, dim=32)
hits = index.query(query, top_k=3)
assert len(index) == len(dataset)
assert len(query) == 32
assert len(hits) == 3
assert hits[0].similarity >= hits[-1].similarity
assert cosine_similarity(query, query) > 0.999
def test_retriever_modes_and_same_state_filter(tmp_path: Path) -> None:
dataset = _make_dataset(tmp_path)
index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
first = dataset[0]
retriever = CriticGatedRetriever(index, dim=32)
same_state = retriever.retrieve(
first.observation_inline,
first.instruction,
k=3,
mode="nearest_neighbor",
same_state_group_id=first.group_id,
)
success_only = retriever.retrieve(first.observation_inline, first.instruction, k=3, mode="success_only")
contrastive = retriever.retrieve(
first.observation_inline,
first.instruction,
k=3,
mode="success_failure_contrastive",
)
assert same_state.examples
assert all(example.item.group_id == first.group_id for example in same_state.examples)
assert all(example.item.success for example in success_only.examples)
assert contrastive.examples
assert any(example.role == "positive_successful" for example in contrastive.examples)
def test_critic_gated_retrieval_uses_optional_critic(tmp_path: Path) -> None:
dataset = _make_dataset(tmp_path)
index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
first = dataset[0]
class MockCritic:
def score_atom(self, atom, selected_atoms, context):
del selected_atoms, context
return 10.0 if atom.reward_summary.get("success", 0.0) else 0.0
retriever = CriticGatedRetriever(
index,
critic=MockCritic(),
transfer_context=TransferContext(benchmark_name="CausalStress"),
dim=32,
)
result = retriever.retrieve(first.observation_inline, first.instruction, k=3, mode="critic_gated")
assert result.examples
assert result.examples[0].gate_score >= result.examples[-1].gate_score
def test_retrieval_conditioned_policy_wrapper(tmp_path: Path) -> None:
dataset = _make_dataset(tmp_path)
index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
first = dataset[0]
retriever = CriticGatedRetriever(index, dim=32)
class DummyPolicy:
def forward_policy(self, observation, instruction, retrieved_examples=None):
del observation, instruction
return {"retrieved": len(retrieved_examples or [])}
wrapper = RetrievalConditionedPolicyWrapper(DummyPolicy(), retriever, k=2)
output = wrapper.policy(first.observation_inline, first.instruction)
assert output["retrieved"] == 2
assert len(wrapper.last_retrieved_examples) == 2
prompt = build_retrieval_prompt(first.instruction, wrapper.last_retrieved_examples)
assert "Retrieved exemplars" in prompt
def test_retrieval_eval_baselines(tmp_path: Path) -> None:
dataset = _make_dataset(tmp_path)
index = CILRetrievalIndex.from_dataset(tmp_path, dim=32)
retriever = CriticGatedRetriever(index, dim=32)
queries = [
RetrievalEvalQuery(
observation=record.observation_inline or {},
instruction=record.instruction,
group_id=record.group_id,
)
for record in dataset.records[:2]
]
report = evaluate_retrieval_baselines(retriever, queries, k=3)
assert set(report) == {
"no_retrieval",
"nearest_neighbor",
"success_only",
"success_failure_contrastive",
"critic_gated",
}
assert report["nearest_neighbor"]["retrieval_coverage"] == 1.0
|