EmergentTwilight/MMCoT / experiments /tests /test_thinkmorph_cache.py
EmergentTwilight's picture
download
raw
5.23 kB
import pytest
from PIL import Image
from thinkmorph_sft.prepare_cache import _mode, canonical_source_row
from thinkmorph_sft.reuse_split_view import reuse_split_view
from zebra_sft.full_data import (
_letterbox_square,
_understanding_grid16_images,
validate_row,
)
from zebra_sft.protocol import SYSTEM_PROMPT
from zebra_sft.trajectory import build_trajectory, render_messages
def _raw_row():
image = {"bytes": b"image", "path": None}
return {
"pid": "sample-1",
"question": "Where is the target?",
"answer": "left",
"problem_image_0": image,
"resoning_thought_0": "Inspect the scene.",
"reasoning_image_0": image,
"resoning_thought_1": "The target is on the left.",
"task": "Visual_Search",
}
def test_thinkmorph_row_uses_canonical_v2_event_boundaries():
source = canonical_source_row(_raw_row())
assert source["Question"] == (
"Where is the target?\n<image_start>[problem_image_0]<image_end>"
)
validated = validate_row(
config="Visual_Search",
row=source,
test_fraction=0.0,
validation_fraction=0.0,
)
assert validated["split"] == "train"
trajectory = build_trajectory(
sample_id="sample-1",
config="Visual_Search",
source_shard="fixture.parquet",
source_row=0,
question=source["Question"],
reasoning_trace=source["Text Reasoning Trace"],
final_answer=source["Final Answer"],
image_paths={"problem_image_0": "p.png", "reasoning_image_0": "r.png"},
image_size=256,
)
messages, roles = render_messages(trajectory, lambda path: path)
assistant = messages[2]["content"]
assert trajectory["system"] == SYSTEM_PROMPT
assert roles == ["observation", "generated_state"]
assert assistant[0]["text"] == "<think>\nInspect the scene.\n</think>"
assert assistant[1]["text"] == "\n<image_gen> generate image 8 8"
assert assistant[3]["text"] == "\n<think>\nThe target is on the left.\n</think>"
assert assistant[4]["text"] == "\nFinal answer: left"
def test_cache_mode_treats_empty_target_as_fresh(tmp_path):
target = tmp_path / "cache"
target.mkdir()
assert _mode(target, auto_resume=True) == "fresh"
(target / "config.json").write_text("{}\n")
assert _mode(target, auto_resume=True) == "resume"
with pytest.raises(FileExistsError):
_mode(target, auto_resume=False)
(target / "READY").write_text("complete\n")
assert _mode(target, auto_resume=True) == "complete"
def test_grid16_profile_preserves_input_and_centers_target_on_gray():
observation = Image.new("RGB", (300, 200), (1, 2, 3))
target = Image.new("RGB", (200, 100), (10, 20, 30))
prepared = _understanding_grid16_images(
[observation, target],
["problem_image_0", "reasoning_image_0"],
["reasoning_image_0"],
target_canvas_size=512,
target_background=128,
)
assert prepared[0].size == observation.size
assert prepared[0].tobytes() == observation.tobytes()
assert prepared[1].size == (512, 512)
assert prepared[1].getpixel((0, 0)) == (128, 128, 128)
assert prepared[1].getpixel((256, 256)) == (10, 20, 30)
assert prepared[1].getpixel((0, 127)) == (128, 128, 128)
assert prepared[1].getpixel((0, 128)) == (10, 20, 30)
def test_reuse_split_view_preserves_assignments_and_repoints_shards(tmp_path):
source = tmp_path / "source"
assignment = tmp_path / "assignment"
output = tmp_path / "view"
source.mkdir()
assignment.mkdir()
rows = [
{
"sample_id": f"sample-{index}",
"split": "train",
"group_id": f"source-group-{index}",
"split_score": 0.5,
"config": "Task",
"source": {"shard": "raw.parquet", "row": index, "pid": str(index)},
"cache_shard": f"shards/part-{index}.pt",
}
for index in range(2)
]
assignments = [
{
**row,
"split": "train" if index == 0 else "validation",
"group_id": f"fixed-group-{index}",
"split_score": index / 2,
}
for index, row in enumerate(rows)
]
for root, values in ((source, rows), (assignment, assignments)):
(root / "manifest.jsonl").write_text(
"".join(__import__("json").dumps(row) + "\n" for row in values)
)
(root / "READY").write_text("complete\n")
(source / "config.json").write_text("{}\n")
(assignment / "image_groups.jsonl").write_text("{}\n")
report = reuse_split_view(
source,
assignment,
output,
expected_task_counts={"Task": 2},
expected_split_counts={"train": 1, "validation": 1},
)
projected = [
__import__("json").loads(line)
for line in (output / "manifest.jsonl").read_text().splitlines()
]
assert report["assignment_preserved"] is True
assert [row["split"] for row in projected] == ["train", "validation"]
assert [row["group_id"] for row in projected] == [
"fixed-group-0",
"fixed-group-1",
]
assert all(str(source) in row["cache_shard"] for row in projected)

Xet Storage Details

Size:
5.23 kB
·
Xet hash:
8af3a4741f10c483dac11aee3bed1f4ae600bf1497e4d272ff1670bc37a603ec

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.