figment / tests /test_finetune_v8_data_plan.py
ThomsenDrake's picture
Sync full submission repo state
94cbe85 verified
Raw
History Blame Contribute Delete
4.86 kB
import json
from collections import Counter
def _accepted_v8_row():
from figment.retrieval import load_protocol_cards
from scripts.generate_finetune_data import assemble_teacher_navigator_output
from scripts.generate_finetune_data import build_sft_row
from scripts.generate_finetune_data import case_spec_record
from scripts.generate_finetune_data import generate_case_spec
from scripts.generate_finetune_data import prepare_case
from scripts.generate_finetune_data import score_candidate
cards_by_id = {str(card["card_id"]): card for card in load_protocol_cards()}
spec = generate_case_spec(0, cards_by_id, dataset_version="figment_sft_v8_delta")
prepared = prepare_case(spec, cards_by_id)
candidate = assemble_teacher_navigator_output(
prepared,
{
"facts": ["postpartum fever confirmed", "temperature elevated"],
"missing": ["pregnancy or postpartum status", "bleeding report", "temperature if available"],
"observe": ["temperature if available", "bleeding report", "mental status"],
"checklist": ["cite fever and pregnancy cards"],
"uncertain": ["blood pressure pending"],
"sbar": {
"situation": "postpartum fever",
"background": "field workflow setting",
"assessment_observations_only": "fever and postpartum context confirmed",
"handoff_request": "request protocol review",
},
"script": "I am checking fever and postpartum observations.",
},
)
result = score_candidate(candidate, prepared)
assert result.passed is True, result.reward_components
row = build_sft_row(
prepared=prepared,
result=result,
teacher_model_id="teacher-test",
candidate_total=1,
candidate_passed=1,
)
return row, case_spec_record(prepared)
def test_v8_failure_cycle_targets_multirule_observation_ownership():
from scripts.generate_finetune_data import V8_NAVIGATOR_COUNTS
from scripts.generate_finetune_data import _failure_class_for_index
categories = Counter(
_failure_class_for_index(index, dataset_version="figment_sft_v8_delta")
for index in range(sum(V8_NAVIGATOR_COUNTS.values()))
)
first_twelve = {
_failure_class_for_index(index, dataset_version="figment_sft_v8_delta")
for index in range(12)
}
assert categories == V8_NAVIGATOR_COUNTS
assert {"multi_rule_observation_ownership", "multi_rule_candidate_focus"} <= first_twelve
def test_v8_sft_row_requires_all_fired_clinical_observation_ids():
row, spec_record = _accepted_v8_row()
output = json.loads(row["messages"][1]["content"])
metadata = row["metadata"]
selected_ids = set(output["selected_required_observation_ids"])
assert row["version"] == "figment_sft_v8_delta"
assert spec_record["dataset_version"] == "figment_sft_v8_delta"
assert spec_record["target_protocol_card_id"] == "FEVER-RED-FLAGS-v1"
assert {"PREG-001", "FEVER-001"} <= set(spec_record["expected_red_flag_rule_ids"])
assert "PREG-DANGER-SIGNS-v1" in output["source_cards"]
candidate_ids = [item["card_id"] for item in output["candidate_protocol_pathways"]]
assert candidate_ids == ["FEVER-RED-FLAGS-v1", "PREG-DANGER-SIGNS-v1"]
assert any(item.startswith("FEVER-RED-FLAGS-v1::required_observation::") for item in selected_ids)
assert any(item.startswith("PREG-DANGER-SIGNS-v1::required_observation::") for item in selected_ids)
assert set(metadata["must_include_selected_required_observation_ids"]) <= selected_ids
observation_text = json.dumps(
output["missing_info_to_collect"] + output["next_observations_to_collect"]
).lower()
for cue in ("temperature if available", "pregnancy or postpartum status", "bleeding report", "fever report"):
assert cue in observation_text
assert metadata["v7_training_policy_version"] == 1
def test_v8_full_corpus_wrapper_pins_delta_defaults():
from scripts.generate_v8_full_corpus import DEFAULT_NAVIGATOR_COUNT
from scripts.generate_v8_full_corpus import DEFAULT_OUTPUT_VERSION
from scripts.generate_v8_full_corpus import DEFAULT_REPAIR_COUNT
from scripts.generate_v8_full_corpus import DEFAULT_TEACHER_MODEL_ID
from scripts.generate_v8_full_corpus import build_corpus_args
assert DEFAULT_OUTPUT_VERSION == "figment_sft_v8_delta"
assert DEFAULT_TEACHER_MODEL_ID == "nvidia/nemotron-3-ultra-550b-a55b:free"
assert DEFAULT_NAVIGATOR_COUNT == 400
assert DEFAULT_REPAIR_COUNT == 0
args = build_corpus_args(["--navigator-count", "8", "--repair-count", "0", "--dry-run"])
assert args[args.index("--navigator-count") + 1] == "8"
assert args[args.index("--repair-count") + 1] == "0"
assert args[-1] == "--dry-run"