File size: 2,863 Bytes
2ff4308 | 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 | import json
from stego_olmoe.data import (
DataCollatorForRouterSFT,
JsonlSFTDataset,
TokenizedSFTDataset,
render_messages_prefix_answer,
write_tokenized_cache,
)
class TinyTokenizer:
eos_token = "<eos>"
eos_token_id = 0
pad_token = "<pad>"
pad_token_id = 0
def __len__(self):
return 512
def __call__(self, text, truncation=True, max_length=512, padding=False, add_special_tokens=True, return_tensors=None):
ids = [ord(ch) % 255 + 1 for ch in text]
if add_special_tokens:
ids = [self.eos_token_id] + ids
ids = ids[:max_length]
return {"input_ids": ids, "attention_mask": [1] * len(ids)}
def test_prompt_answer_router_mask_answer_only(tmp_path):
path = tmp_path / "data.jsonl"
path.write_text(json.dumps({"prompt": "A", "answer": "B"}) + "\n", encoding="utf-8")
dataset = JsonlSFTDataset(path, TinyTokenizer(), max_length=64)
item = dataset[0]
assert item["router_loss_mask"].sum().item() >= 1
assert item["router_loss_mask"][0].item() == 0
def test_collator_pads_router_mask(tmp_path):
path = tmp_path / "data.jsonl"
path.write_text(
json.dumps({"text": "short"}) + "\n" + json.dumps({"text": "a little longer"}) + "\n",
encoding="utf-8",
)
tokenizer = TinyTokenizer()
dataset = JsonlSFTDataset(path, tokenizer, max_length=64)
batch = DataCollatorForRouterSFT(tokenizer)([dataset[0], dataset[1]])
assert batch["input_ids"].shape[0] == 2
assert batch["router_loss_mask"].shape == batch["attention_mask"].shape
def test_messages_render_last_assistant_target():
prefix, answer = render_messages_prefix_answer(
[
{"role": "system", "content": "Be concise."},
{"role": "user", "content": "Question"},
{"role": "assistant", "content": "First answer"},
{"role": "user", "content": "Follow-up"},
{"role": "assistant", "content": "Final answer"},
]
)
assert prefix.endswith("Assistant:")
assert "First answer" in prefix
assert "Final answer" not in prefix
assert answer == "Final answer"
def test_tokenized_cache_round_trips(tmp_path):
jsonl_path = tmp_path / "data.jsonl"
cache_path = tmp_path / "data.pt"
jsonl_path.write_text(
json.dumps({"prompt": "A", "answer": "B"}) + "\n" + json.dumps({"text": "free text"}) + "\n",
encoding="utf-8",
)
metadata = write_tokenized_cache(jsonl_path, cache_path, TinyTokenizer(), max_length=64)
jsonl_dataset = JsonlSFTDataset(jsonl_path, TinyTokenizer(), max_length=64)
cache_dataset = TokenizedSFTDataset(cache_path)
assert metadata["n_records"] == 2
assert len(cache_dataset) == len(jsonl_dataset)
for key, value in jsonl_dataset[0].items():
assert value.equal(cache_dataset[0][key])
|