From 25d4ea3f91def58aaec2ae429f376087b19ba194 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 19 Jun 2026 14:53:35 +0800 Subject: [PATCH] =?UTF-8?q?refactor=20:=20=E5=8E=8B=E7=BC=A9=E6=B5=8B?= =?UTF-8?q?=E8=AF=95=E4=BB=A3=E7=A0=81=EF=BC=8C=E6=B6=88=E9=99=A4=E9=87=8D?= =?UTF-8?q?=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fixture 替代重复实例化和 tokenizer 落盘 - parametrize 合并同构测试 - helper 消除 save_h5 + DatasetFactory.load 样板 - 净减 272 行 --- tests/data/conftest.py | 33 ++++ tests/data/test_dataset.py | 180 +++++++-------------- tests/data/test_preprocess_builder.py | 113 +++++-------- tests/data/test_preprocess_pipeline.py | 101 +----------- tests/inference/test_tool_parser.py | 209 ++++++++----------------- tests/module/test_encoder.py | 96 ++++-------- 6 files changed, 230 insertions(+), 502 deletions(-) diff --git a/tests/data/conftest.py b/tests/data/conftest.py index 5d5fe05..cfb25c7 100644 --- a/tests/data/conftest.py +++ b/tests/data/conftest.py @@ -1,3 +1,5 @@ +import json +import os import tempfile import pytest @@ -8,6 +10,7 @@ from astrai.config.preprocess_config import ( PipelineConfig, ProcessingConfig, ) +from astrai.preprocessing.builder import SectionedMaskBuilder from astrai.tokenize import AutoTokenizer _SPECIAL_TOKENS_CONFIG = { @@ -200,3 +203,33 @@ def make_grpo_no_template_config(): mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=2048), ) + + +@pytest.fixture +def builder(): + return SectionedMaskBuilder() + + +@pytest.fixture +def tokenizer_dir(temp_dir, test_tokenizer): + d = os.path.join(temp_dir, "tok") + os.makedirs(d, exist_ok=True) + test_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json")) + with open(os.path.join(d, "tokenizer_config.json"), "w") as f: + json.dump( + {"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}}, f + ) + return d + + +@pytest.fixture +def chat_tokenizer_dir(temp_dir, chat_tokenizer): + d = os.path.join(temp_dir, "tok") + os.makedirs(d, exist_ok=True) + chat_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json")) + with open(os.path.join(d, "tokenizer_config.json"), "w") as f: + json.dump( + {"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE}, + f, + ) + return d diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index c0f6ba5..6f800dd 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -15,28 +15,34 @@ from astrai.dataset.storage import ( ) +def _rand_seq(length, vocab=1000): + return torch.randint(0, vocab, (length,), dtype=torch.int64) + + +def _make_seq_dataset( + test_dir, name="data", seq_length=200, train_type="seq", data=None, **load_kwargs +): + if data is None: + data = {"sequence": [_rand_seq(seq_length)]} + save_h5(test_dir, name, data) + return DatasetFactory.load( + train_type, + test_dir, + window_size=load_kwargs.pop("window_size", 64), + **load_kwargs, + ) + + def test_dataset_loader_random_paths(base_test_env): """Test dataset loader with multiple random paths""" test_dir = base_test_env["test_dir"] - # Create multiple mmap dataset directories with random data num_files = np.random.randint(2, 5) - for i in range(num_files): seq_length = np.random.randint(200, 400) - dummy_data = { - "sequence": [ - torch.randint(0, 1000, (seq_length,), dtype=torch.int64) - for _ in range(10) - ], - } - save_h5(test_dir, f"data_{i}", dummy_data) - - # Test loading with multiple paths - loaded_dataset = DatasetFactory.load( - train_type="seq", - load_path=test_dir, - window_size=64, + dummy_data = {"sequence": [_rand_seq(seq_length) for _ in range(10)]} + loaded_dataset = _make_seq_dataset( + test_dir, f"data_{i}", seq_length, data=dummy_data ) assert loaded_dataset is not None assert len(loaded_dataset) > 0 @@ -54,23 +60,15 @@ def test_dpo_strategy_with_random_data(base_test_env): """Test DPO strategy with randomized preference data""" test_dir = base_test_env["test_dir"] - # Create DPO-style data with memory mapping format seq_length = np.random.randint(100, 200) - dummy_data = { - "chosen": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)], - "rejected": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)], + "chosen": [_rand_seq(seq_length)], + "rejected": [_rand_seq(seq_length)], "chosen_mask": [torch.ones(seq_length, dtype=torch.bool)], "rejected_mask": [torch.ones(seq_length, dtype=torch.bool)], } - - save_h5(test_dir, "dpo_data", dummy_data) - - # Load DPO dataset - dpo_dataset = DatasetFactory.load( - train_type="dpo", - load_path=test_dir, - window_size=64, + dpo_dataset = _make_seq_dataset( + test_dir, "dpo_data", seq_length, train_type="dpo", data=dummy_data ) assert dpo_dataset is not None @@ -92,22 +90,14 @@ def test_sft_dataset_with_random_data(base_test_env): """Test SFT dataset with random data""" test_dir = base_test_env["test_dir"] - # Create SFT-style data with memory mapping format seq_length = np.random.randint(100, 200) - dummy_data = { - "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)], + "sequence": [_rand_seq(seq_length)], "loss_mask": [torch.ones(seq_length, dtype=torch.bool)], "position_ids": [torch.arange(seq_length, dtype=torch.int32)], } - - save_h5(test_dir, "sft_data", dummy_data) - - # Load SFT dataset - sft_dataset = DatasetFactory.load( - train_type="sft", - load_path=test_dir, - window_size=64, + sft_dataset = _make_seq_dataset( + test_dir, "sft_data", seq_length, train_type="sft", data=dummy_data ) assert sft_dataset is not None @@ -128,25 +118,11 @@ def test_dataset_with_custom_stride(base_test_env): """Test dataset with custom stride parameter""" test_dir = base_test_env["test_dir"] - # Create test data - seq_length = 200 - dummy_data = { - "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)], - } - - save_h5(test_dir, "stride_test_data", dummy_data) - - # Test with custom stride custom_stride = 32 - dataset = DatasetFactory.load( - train_type="seq", load_path=test_dir, window_size=64, stride=custom_stride - ) - + dataset = _make_seq_dataset(test_dir, "stride_test_data", stride=custom_stride) assert dataset is not None assert len(dataset) > 0 - # With stride 32 and window 64 on 200 length data, we should get more samples - # than with default stride (which equals window size) default_stride_dataset = DatasetFactory.load( train_type="seq", load_path=test_dir, @@ -157,25 +133,11 @@ def test_dataset_with_custom_stride(base_test_env): def test_dataset_count_property(base_test_env): - """Test the count property returns correct raw token count""" test_dir = base_test_env["test_dir"] - - seq_length = 200 - dummy_data = { - "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)], - } - - save_h5(test_dir, "count_test_data", dummy_data) - - dataset = DatasetFactory.load( - train_type="seq", - load_path=test_dir, - window_size=64, - ) - - assert dataset.count == seq_length - assert dataset.count > len(dataset) # raw tokens > windows - assert len(dataset) == (seq_length - 1 - 64) // 64 + 1 + dataset = _make_seq_dataset(test_dir, "count_test_data") + assert dataset.count == 200 + assert dataset.count > len(dataset) + assert len(dataset) == (200 - 1 - 64) // 64 + 1 def test_empty_dataset_count(): @@ -186,17 +148,10 @@ def test_empty_dataset_count(): def test_dataset_too_short_for_window(base_test_env): - """Dataset shorter than window_size returns __len__ == 0""" test_dir = base_test_env["test_dir"] - seq_length = 30 - save_h5( - test_dir, - "short", - {"sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)]}, - ) - dataset = DatasetFactory.load("seq", test_dir, window_size=64) + dataset = _make_seq_dataset(test_dir, "short", seq_length=30) assert len(dataset) == 0 - assert dataset.count == seq_length + assert dataset.count == 30 def test_unloaded_dataset_getitem_raises(): @@ -220,12 +175,8 @@ def test_store_unloaded_len(): def test_store_fetch_begin_equals_end(base_test_env): - """Store.fetch with begin == end returns empty tensor""" test_dir = base_test_env["test_dir"] - dummy = {"sequence": [torch.randint(0, 1000, (100,), dtype=torch.int64)]} - save_h5(test_dir, "empty_fetch", dummy) - - dataset = DatasetFactory.load("seq", test_dir, window_size=32) + dataset = _make_seq_dataset(test_dir, "empty_fetch", seq_length=100, window_size=32) result = dataset.storage.fetch(10, 10, "sequence") assert result.numel() == 0 @@ -299,12 +250,8 @@ def test_save_load_bin_roundtrip(base_test_env): def test_mmap_store_load_and_fetch(base_test_env): - """MmapStore loads bin data and fetches correctly""" test_dir = base_test_env["test_dir"] - - data = { - "sequence": [torch.randint(0, 1000, (200,), dtype=torch.int64)], - } + data = {"sequence": [_rand_seq(200)]} save_bin(test_dir, data) store = StoreFactory.create("bin") @@ -317,14 +264,9 @@ def test_mmap_store_load_and_fetch(base_test_env): def test_mmap_dataset_load(base_test_env): - """DatasetFactory.load auto-detects bin format""" test_dir = base_test_env["test_dir"] - - data = { - "sequence": [torch.randint(0, 1000, (200,), dtype=torch.int64)], - } + data = {"sequence": [_rand_seq(200)]} save_bin(test_dir, data) - dataset = DatasetFactory.load("seq", test_dir, window_size=64) assert len(dataset) > 0 assert dataset.count == 200 @@ -348,19 +290,16 @@ def test_normalize_mixed_empty_key(): def test_grpo_dataset_dtype(base_test_env): - """GRPODataset returns correct dtypes""" test_dir = base_test_env["test_dir"] - - seq_len = 100 - data = { - "prompts": [torch.randint(0, 100, (seq_len,), dtype=torch.int32)], - "responses": [torch.randint(0, 100, (seq_len,), dtype=torch.int32)], - "masks": [torch.ones(seq_len, dtype=torch.int32)], - "rewards": [torch.ones(seq_len, dtype=torch.float32)], + dummy_data = { + "prompts": [torch.randint(0, 100, (100,), dtype=torch.int32)], + "responses": [torch.randint(0, 100, (100,), dtype=torch.int32)], + "masks": [torch.ones(100, dtype=torch.int32)], + "rewards": [torch.ones(100, dtype=torch.float32)], } - save_h5(test_dir, "grpo_dtype", data) - - dataset = DatasetFactory.load("grpo", test_dir, window_size=32) + dataset = _make_seq_dataset( + test_dir, "grpo_dtype", train_type="grpo", data=dummy_data, window_size=32 + ) item = dataset[0] assert item["prompts"].dtype == torch.long @@ -370,18 +309,16 @@ def test_grpo_dataset_dtype(base_test_env): def test_grpo_dataset_load(base_test_env): - """GRPODataset loads and returns correct keys""" test_dir = base_test_env["test_dir"] - seq_len = 200 - data = { - "prompts": [torch.randint(0, 1000, (seq_len,), dtype=torch.int64)], - "responses": [torch.randint(0, 1000, (seq_len,), dtype=torch.int64)], - "masks": [torch.ones(seq_len, dtype=torch.int64)], - "rewards": [torch.rand(seq_len, dtype=torch.float32)], + dummy_data = { + "prompts": [_rand_seq(200)], + "responses": [_rand_seq(200)], + "masks": [torch.ones(200, dtype=torch.int64)], + "rewards": [torch.rand(200, dtype=torch.float32)], } - save_h5(test_dir, "grpo_test", data) - - dataset = DatasetFactory.load("grpo", test_dir, window_size=64) + dataset = _make_seq_dataset( + test_dir, "grpo_test", train_type="grpo", data=dummy_data + ) assert len(dataset) > 0 item = dataset[0] assert "prompts" in item @@ -400,7 +337,6 @@ def test_detect_format_bin_dir(base_test_env): def test_store_fetch_multi_key(base_test_env): - """Store.fetch with List[str] returns Dict[str, Tensor]""" test_dir = base_test_env["test_dir"] save_h5( test_dir, @@ -410,7 +346,6 @@ def test_store_fetch_multi_key(base_test_env): "loss_mask": [torch.ones(100, dtype=torch.int64)], }, ) - store = StoreFactory.create("h5") store.load(test_dir) result = store.fetch(10, 20, ["sequence", "loss_mask"]) @@ -420,10 +355,8 @@ def test_store_fetch_multi_key(base_test_env): def test_store_fetch_out_of_bounds(base_test_env): - """Store.fetch raises ValueError for out-of-bounds indices""" test_dir = base_test_env["test_dir"] save_h5(test_dir, "bounds", {"sequence": [torch.randint(0, 100, (50,))]}) - store = StoreFactory.create("h5") store.load(test_dir) with pytest.raises(ValueError, match="out of bounds"): @@ -435,10 +368,7 @@ def test_store_fetch_out_of_bounds(base_test_env): def test_dataset_load_explicit_storage_type(base_test_env): - """DatasetFactory.load with explicit storage_type bypasses auto-detect""" test_dir = base_test_env["test_dir"] - save_h5(test_dir, "explicit", {"sequence": [torch.randint(0, 100, (200,))]}) - - dataset = DatasetFactory.load("seq", test_dir, window_size=64, storage_type="h5") + dataset = _make_seq_dataset(test_dir, "explicit", storage_type="h5") assert len(dataset) > 0 assert dataset.count == 200 diff --git a/tests/data/test_preprocess_builder.py b/tests/data/test_preprocess_builder.py index eb619ab..7f886f0 100644 --- a/tests/data/test_preprocess_builder.py +++ b/tests/data/test_preprocess_builder.py @@ -1,3 +1,5 @@ +import pytest + from astrai.config.preprocess_config import ( InputConfig, OutputConfig, @@ -20,9 +22,8 @@ from tests.data.conftest import ( ) -def test_chat_simple(chat_tokenizer): +def test_chat_simple(chat_tokenizer, builder): config = make_chat_config() - builder = SectionedMaskBuilder() item = { "messages": [ {"role": "system", "content": "You are helpful."}, @@ -46,9 +47,8 @@ def test_chat_simple(chat_tokenizer): assert trained < total -def test_chat_mask_only_assistant(chat_tokenizer): +def test_chat_mask_only_assistant(chat_tokenizer, builder): config = make_chat_config() - builder = SectionedMaskBuilder() item = { "messages": [ {"role": "user", "content": "What is 2+2?"}, @@ -66,14 +66,22 @@ def test_chat_mask_only_assistant(chat_tokenizer): assert len(masked) > 0 -def test_chat_all_masked(chat_tokenizer): +@pytest.mark.parametrize( + "mask_rules,mask_default,expect_nonzero", + [ + ({"system": "mask", "user": "mask", "assistant": "mask"}, "mask", False), + ({}, "train", True), + ], +) +def test_chat_uniform_masking( + mask_rules, mask_default, expect_nonzero, chat_tokenizer, builder +): config = PipelineConfig( input=InputConfig(sections=_CHAT_SECTIONS), - mask={"system": "mask", "user": "mask", "assistant": "mask"}, - mask_default="mask", + mask=mask_rules, + mask_default=mask_default, preprocessing=ProcessingConfig(max_seq_len=2048), ) - builder = SectionedMaskBuilder() item = { "messages": [ {"role": "system", "content": "You are helpful."}, @@ -81,35 +89,20 @@ def test_chat_all_masked(chat_tokenizer): ] } result = builder.build(item, config, chat_tokenizer) - assert sum(result["loss_mask"]) == 0 + masked_count = sum(result["loss_mask"]) + if expect_nonzero: + assert masked_count > 0 + else: + assert masked_count == 0 -def test_chat_all_trained(chat_tokenizer): - config = PipelineConfig( - input=InputConfig(sections=_CHAT_SECTIONS), - mask={}, - mask_default="train", - preprocessing=ProcessingConfig(max_seq_len=2048), - ) - builder = SectionedMaskBuilder() - item = { - "messages": [ - {"role": "system", "content": "You are helpful."}, - {"role": "assistant", "content": "Hi there!"}, - ] - } - result = builder.build(item, config, chat_tokenizer) - assert sum(result["loss_mask"]) == len(result["sequence"]) - 1 - - -def test_chat_empty_messages(chat_tokenizer): +def test_chat_empty_messages(chat_tokenizer, builder): config = make_chat_config() - builder = SectionedMaskBuilder() assert builder.build({"messages": []}, config, chat_tokenizer) is None assert builder.build({}, config, chat_tokenizer) is None -def test_chat_domain_extraction(chat_tokenizer): +def test_chat_domain_extraction(chat_tokenizer, builder): config = PipelineConfig( input=InputConfig(sections=_CHAT_SECTIONS), mask={"assistant": "train"}, @@ -117,7 +110,6 @@ def test_chat_domain_extraction(chat_tokenizer): preprocessing=ProcessingConfig(max_seq_len=2048), output=OutputConfig(domain_key="source"), ) - builder = SectionedMaskBuilder() item = { "messages": [ {"role": "user", "content": "Hi"}, @@ -129,14 +121,13 @@ def test_chat_domain_extraction(chat_tokenizer): assert result["domain"] == "wiki" -def test_chat_truncation(chat_tokenizer): +def test_chat_truncation(chat_tokenizer, builder): config = PipelineConfig( input=InputConfig(sections=_CHAT_SECTIONS), mask={"assistant": "train"}, mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=10), ) - builder = SectionedMaskBuilder() item = { "messages": [ { @@ -151,18 +142,16 @@ def test_chat_truncation(chat_tokenizer): assert len(result["loss_mask"]) == len(result["sequence"]) -def test_instruction_basic(test_tokenizer): +def test_instruction_basic(test_tokenizer, builder): config = make_instruction_config() - builder = SectionedMaskBuilder() item = {"prompt": "Translate to French: Hello", "response": "Bonjour"} result = builder.build(item, config, test_tokenizer) assert result is not None assert len(result["sequence"]) == len(result["loss_mask"]) -def test_instruction_prompt_masked(test_tokenizer): +def test_instruction_prompt_masked(test_tokenizer, builder): config = make_instruction_config() - builder = SectionedMaskBuilder() item = {"prompt": "hello", "response": "world"} result = builder.build(item, config, test_tokenizer) mask = result["loss_mask"] @@ -175,7 +164,7 @@ def test_instruction_prompt_masked(test_tokenizer): assert all(m == 1 for m in mask[p_len:]) -def test_instruction_train_on_prompt(test_tokenizer): +def test_instruction_train_on_prompt(test_tokenizer, builder): config = PipelineConfig( input=InputConfig( sections=[ @@ -185,7 +174,6 @@ def test_instruction_train_on_prompt(test_tokenizer): ), preprocessing=ProcessingConfig(max_seq_len=2048), ) - builder = SectionedMaskBuilder() item = {"prompt": "hello", "response": "world"} result = builder.build(item, config, test_tokenizer) mask = result["loss_mask"] @@ -196,9 +184,8 @@ def test_instruction_train_on_prompt(test_tokenizer): assert all(m == 1 for m in mask[:p_len]) -def test_text_basic(test_tokenizer): +def test_text_basic(test_tokenizer, builder): config = make_text_config() - builder = SectionedMaskBuilder() item = {"text": "Hello world. This is a test document."} result = builder.build(item, config, test_tokenizer) assert result is not None @@ -207,41 +194,37 @@ def test_text_basic(test_tokenizer): assert "loss_mask" not in result -def test_text_empty(test_tokenizer): +def test_text_empty(test_tokenizer, builder): config = make_text_config() - builder = SectionedMaskBuilder() assert builder.build({"text": ""}, config, test_tokenizer) is None assert builder.build({"text": " "}, config, test_tokenizer) is None -def test_text_too_short(test_tokenizer): +def test_text_too_short(test_tokenizer, builder): config = PipelineConfig( input=InputConfig(sections=_TEXT_SECTIONS), preprocessing=ProcessingConfig(min_chars=100), ) - builder = SectionedMaskBuilder() assert builder.build({"text": "short"}, config, test_tokenizer) is None -def test_text_truncation(test_tokenizer): +def test_text_truncation(test_tokenizer, builder): config = PipelineConfig( input=InputConfig(sections=_TEXT_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1), ) - builder = SectionedMaskBuilder() item = {"text": "This is a very long text that should be truncated"} result = builder.build(item, config, test_tokenizer) assert len(result["sequence"]) <= 3 -def test_sectioned_chat(chat_tokenizer): +def test_sectioned_chat(chat_tokenizer, builder): config = PipelineConfig( input=InputConfig(sections=_CHAT_SECTIONS), mask={"system": "mask", "user": "mask", "assistant": "train"}, mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=2048), ) - builder = SectionedMaskBuilder() item = { "messages": [ {"role": "user", "content": "What is 2+2?"}, @@ -255,12 +238,11 @@ def test_sectioned_chat(chat_tokenizer): assert 0 in result["loss_mask"] -def test_sectioned_instruction(test_tokenizer): +def test_sectioned_instruction(test_tokenizer, builder): config = PipelineConfig( input=InputConfig(sections=_INSTRUCTION_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0), ) - builder = SectionedMaskBuilder() item = {"prompt": "Q: Why?", "response": "A: Because."} result = builder.build(item, config, test_tokenizer) assert result is not None @@ -269,24 +251,22 @@ def test_sectioned_instruction(test_tokenizer): assert mask[-1] == 1 -def test_sectioned_text(test_tokenizer): +def test_sectioned_text(test_tokenizer, builder): config = PipelineConfig( input=InputConfig(sections=_TEXT_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1), ) - builder = SectionedMaskBuilder() item = {"text": "Hello world, this is a test."} result = builder.build(item, config, test_tokenizer) assert result is not None assert "loss_mask" not in result -def test_sectioned_text_too_short(test_tokenizer): +def test_sectioned_text_too_short(test_tokenizer, builder): config = PipelineConfig( input=InputConfig(sections=_TEXT_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100), ) - builder = SectionedMaskBuilder() assert builder.build({"text": "short"}, config, test_tokenizer) is None @@ -296,13 +276,12 @@ def test_factory_registered(): def test_factory_create(): - builder = MaskBuilderFactory.create("sectioned") - assert isinstance(builder, SectionedMaskBuilder) + builder_obj = MaskBuilderFactory.create("sectioned") + assert isinstance(builder_obj, SectionedMaskBuilder) -def test_dpo_chat_basic(chat_tokenizer): +def test_dpo_chat_basic(chat_tokenizer, builder): config = make_dpo_chat_config() - builder = SectionedMaskBuilder() item = { "chosen": [ {"role": "user", "content": "What is 2+2?"}, @@ -319,16 +298,14 @@ def test_dpo_chat_basic(chat_tokenizer): assert "rejected" in result assert "chosen_mask" in result assert "rejected_mask" in result - assert "domain" in result assert len(result["chosen"]) == len(result["chosen_mask"]) assert len(result["rejected"]) == len(result["rejected_mask"]) assert sum(result["chosen_mask"]) > 0 assert sum(result["rejected_mask"]) > 0 -def test_dpo_chosen_only_trained(chat_tokenizer): +def test_dpo_chosen_only_trained(chat_tokenizer, builder): config = make_dpo_chat_config() - builder = SectionedMaskBuilder() item = { "chosen": [ {"role": "user", "content": "Hi"}, @@ -346,15 +323,13 @@ def test_dpo_chosen_only_trained(chat_tokenizer): assert 1 in result["rejected_mask"] -def test_dpo_missing_field_is_none(chat_tokenizer): +def test_dpo_missing_field_is_none(chat_tokenizer, builder): config = make_dpo_chat_config() - builder = SectionedMaskBuilder() assert builder.build({"chosen": [], "rejected": []}, config, chat_tokenizer) is None -def test_grpo_basic(chat_tokenizer): +def test_grpo_basic(chat_tokenizer, builder): config = make_grpo_config() - builder = SectionedMaskBuilder() item = { "prompt": [{"role": "user", "content": "What is 2+2?"}], "responses": ["4", "The answer is four", "Four", "2+2=4"], @@ -370,9 +345,8 @@ def test_grpo_basic(chat_tokenizer): assert result["rewards"] == [1.0, 0.5, 0.8, 0.2] -def test_grpo_response_tokens_all_trained(chat_tokenizer): +def test_grpo_response_tokens_all_trained(chat_tokenizer, builder): config = make_grpo_config() - builder = SectionedMaskBuilder() item = { "prompt": [{"role": "user", "content": "Q"}], "responses": ["A", "B"], @@ -384,9 +358,8 @@ def test_grpo_response_tokens_all_trained(chat_tokenizer): assert len(masks) == len(result["responses"]) -def test_grpo_single_reward(chat_tokenizer): +def test_grpo_single_reward(chat_tokenizer, builder): config = make_grpo_config() - builder = SectionedMaskBuilder() item = { "prompt": [{"role": "user", "content": "Q"}], "responses": ["A"], diff --git a/tests/data/test_preprocess_pipeline.py b/tests/data/test_preprocess_pipeline.py index e28e90c..250de4b 100644 --- a/tests/data/test_preprocess_pipeline.py +++ b/tests/data/test_preprocess_pipeline.py @@ -10,9 +10,7 @@ from astrai.config.preprocess_config import ( from astrai.preprocessing.pipeline import Pipeline, filter_by_length from tests.data.conftest import ( _CHAT_SECTIONS, - _CHAT_TEMPLATE, _INSTRUCTION_SECTIONS, - _SPECIAL_TOKENS_CONFIG, _TEXT_SECTIONS, make_dpo_chat_config, make_grpo_no_template_config, @@ -26,19 +24,7 @@ def test_filter_by_length(): assert filter_by_length("just right", min_len=5, max_len=20) -def test_full_chat_pipeline(temp_dir, chat_tokenizer): - tokenizer_dir = os.path.join(temp_dir, "tok") - os.makedirs(tokenizer_dir, exist_ok=True) - chat_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json")) - with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f: - json.dump( - { - "special_tokens": _SPECIAL_TOKENS_CONFIG, - "chat_template": _CHAT_TEMPLATE, - }, - f, - ) - +def test_full_chat_pipeline(temp_dir, chat_tokenizer_dir): jsonl_path = os.path.join(temp_dir, "chat.jsonl") with open(jsonl_path, "w", encoding="utf-8") as f: f.write( @@ -78,7 +64,7 @@ def test_full_chat_pipeline(temp_dir, chat_tokenizer): config=config, input_paths=[jsonl_path], output_dir=out_dir, - tokenizer_path=tokenizer_dir, + tokenizer_path=chat_tokenizer_dir, ).run() meta_path = os.path.join(out_dir, "__default__", "shard_0000", "meta.json") @@ -91,21 +77,7 @@ def test_full_chat_pipeline(temp_dir, chat_tokenizer): assert meta["loss_mask"]["dtype"] == "int32" -def test_full_text_pipeline(temp_dir, test_tokenizer): - tokenizer_dir = os.path.join(temp_dir, "tok") - os.makedirs(tokenizer_dir, exist_ok=True) - test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json")) - with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f: - json.dump( - { - "special_tokens": { - "pad_token": "<|_pad_|>", - "unk_token": "<|_unk_|>", - } - }, - f, - ) - +def test_full_text_pipeline(temp_dir, tokenizer_dir): jsonl_path = os.path.join(temp_dir, "text.jsonl") with open(jsonl_path, "w", encoding="utf-8") as f: f.write( @@ -145,24 +117,9 @@ def test_full_text_pipeline(temp_dir, test_tokenizer): meta = json.load(f) assert "sequence" in meta assert "loss_mask" not in meta - assert meta["sequence"]["dtype"] == "int32" -def test_full_instruction_pipeline(temp_dir, test_tokenizer): - tokenizer_dir = os.path.join(temp_dir, "tok") - os.makedirs(tokenizer_dir, exist_ok=True) - test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json")) - with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f: - json.dump( - { - "special_tokens": { - "pad_token": "<|_pad_|>", - "unk_token": "<|_unk_|>", - } - }, - f, - ) - +def test_full_instruction_pipeline(temp_dir, tokenizer_dir): jsonl_path = os.path.join(temp_dir, "instruct.jsonl") with open(jsonl_path, "w", encoding="utf-8") as f: f.write( @@ -206,25 +163,9 @@ def test_full_instruction_pipeline(temp_dir, test_tokenizer): meta = json.load(f) assert "sequence" in meta assert "loss_mask" in meta - assert meta["sequence"]["dtype"] == "int32" - assert meta["loss_mask"]["dtype"] == "int32" -def test_dtype_override(temp_dir, test_tokenizer): - tokenizer_dir = os.path.join(temp_dir, "tok") - os.makedirs(tokenizer_dir, exist_ok=True) - test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json")) - with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f: - json.dump( - { - "special_tokens": { - "pad_token": "<|_pad_|>", - "unk_token": "<|_unk_|>", - } - }, - f, - ) - +def test_dtype_override(temp_dir, tokenizer_dir): jsonl_path = os.path.join(temp_dir, "data.jsonl") with open(jsonl_path, "w", encoding="utf-8") as f: f.write(json.dumps({"prompt": "Q", "response": "A"}) + "\n") @@ -252,19 +193,7 @@ def test_dtype_override(temp_dir, test_tokenizer): assert meta["loss_mask"]["dtype"] == "bool" -def test_dpo_pipeline(temp_dir, chat_tokenizer): - tokenizer_dir = os.path.join(temp_dir, "tok") - os.makedirs(tokenizer_dir, exist_ok=True) - chat_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json")) - with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f: - json.dump( - { - "special_tokens": _SPECIAL_TOKENS_CONFIG, - "chat_template": _CHAT_TEMPLATE, - }, - f, - ) - +def test_dpo_pipeline(temp_dir, chat_tokenizer_dir): jsonl_path = os.path.join(temp_dir, "dpo.jsonl") with open(jsonl_path, "w", encoding="utf-8") as f: f.write( @@ -288,7 +217,7 @@ def test_dpo_pipeline(temp_dir, chat_tokenizer): config=make_dpo_chat_config(), input_paths=[jsonl_path], output_dir=out_dir, - tokenizer_path=tokenizer_dir, + tokenizer_path=chat_tokenizer_dir, ).run() meta_path = os.path.join(out_dir, "__default__", "shard_0000", "meta.json") @@ -302,21 +231,7 @@ def test_dpo_pipeline(temp_dir, chat_tokenizer): assert "sequence" not in meta -def test_grpo_pipeline(temp_dir, test_tokenizer): - tokenizer_dir = os.path.join(temp_dir, "tok") - os.makedirs(tokenizer_dir, exist_ok=True) - test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json")) - with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f: - json.dump( - { - "special_tokens": { - "pad_token": "<|_pad_|>", - "unk_token": "<|_unk_|>", - } - }, - f, - ) - +def test_grpo_pipeline(temp_dir, tokenizer_dir): jsonl_path = os.path.join(temp_dir, "grpo.jsonl") with open(jsonl_path, "w", encoding="utf-8") as f: f.write( diff --git a/tests/inference/test_tool_parser.py b/tests/inference/test_tool_parser.py index 26d57b8..2a6b945 100644 --- a/tests/inference/test_tool_parser.py +++ b/tests/inference/test_tool_parser.py @@ -13,65 +13,26 @@ from astrai.inference.api.tool_parser import ( ) -def test_scan_complete_simple(): - end, complete = _scan_json('{"key": "value"}', 0) - assert complete is True - assert end == len('{"key": "value"}') - - -def test_scan_complete_nested(): - text = '{"outer": {"inner": 1}}' +@pytest.mark.parametrize( + "text,expected_complete,check_end_eq_len", + [ + ('{"key": "value"}', True, True), + ('{"outer": {"inner": 1}}', True, True), + ('{"key": "value"', False, False), + ('{"outer": {"inner": 1}', False, False), + ('{"key": "a{b}c"} extra', True, False), + (r'{"key": "a\"b"}', True, False), + ('{"a": {"b": {"c": {"d": {"e": 5}}}}}', True, True), + ('{"items": [{"x": 1}, {"x": 2}]}', True, True), + ('{"fn": "function() { return 1; }"}', True, False), + ('{"key": "\u5317\u4eac"}', True, False), + ], +) +def test_scan_json(text, expected_complete, check_end_eq_len): end, complete = _scan_json(text, 0) - assert complete is True - assert end == len(text) - - -def test_scan_incomplete_unclosed(): - end, complete = _scan_json('{"key": "value"', 0) - assert complete is False - - -def test_scan_incomplete_nested(): - end, complete = _scan_json('{"outer": {"inner": 1}', 0) - assert complete is False - - -def test_scan_string_braces_ignored(): - text = '{"key": "a{b}c"} extra' - end, complete = _scan_json(text, 0) - assert complete is True - - -def test_scan_escaped_quote_ignored(): - text = r'{"key": "a\"b"}' - end, complete = _scan_json(text, 0) - assert complete is True - - -def test_scan_deeply_nested(): - text = '{"a": {"b": {"c": {"d": {"e": 5}}}}}' - end, complete = _scan_json(text, 0) - assert complete is True - assert end == len(text) - - -def test_scan_array_with_braces(): - text = '{"items": [{"x": 1}, {"x": 2}]}' - end, complete = _scan_json(text, 0) - assert complete is True - assert end == len(text) - - -def test_scan_code_in_string(): - text = '{"fn": "function() { return 1; }"}' - end, complete = _scan_json(text, 0) - assert complete is True - - -def test_scan_unicode_chars(): - text = '{"key": "\u5317\u4eac"}' - end, complete = _scan_json(text, 0) - assert complete is True + assert complete is expected_complete + if check_end_eq_len: + assert end == len(text) def test_find_single_tool_call(): @@ -141,10 +102,7 @@ def test_find_arguments_with_array(): def test_find_arguments_with_nested_array_of_objects(): - text = ( - '{"name": "batch", ' - '"arguments": {"rows": [{"id": 1, "val": "a"}, {"id": 2, "val": "b"}]}}' - ) + text = '{"name": "batch", "arguments": {"rows": [{"id": 1, "val": "a"}, {"id": 2, "val": "b"}]}}' results = _find_tool_calls(text) assert len(results) == 1 assert '"rows"' in results[0]["args"] @@ -206,38 +164,26 @@ def test_find_extracts_correct_arg_start_position(): assert json_str == text -def test_partial_with_name(): - result = _find_partial_tool_call('{"name": "func", "arguments": {"city"') - assert result is not None - assert result["name"] == "func" - assert result["complete"] is False - - -def test_partial_with_full_args(): - result = _find_partial_tool_call('{"name": "func", "arguments": {"city": "BJ"}}') - assert result is not None - assert result["name"] == "func" - - -def test_partial_no_match(): - assert _find_partial_tool_call("plain text") is None - - -def test_partial_no_name_yet(): - assert _find_partial_tool_call('{"nam') is None - - -def test_partial_deeply_nested(): - result = _find_partial_tool_call('{"name": "deep", "arguments": {"a": {"b": {"c": ') - assert result is not None - assert result["name"] == "deep" - assert '"a"' in result["args"] - - -def test_partial_array_incomplete(): - result = _find_partial_tool_call('{"name": "batch", "arguments": {"items": [1, 2, ') - assert result is not None - assert result["name"] == "batch" +@pytest.mark.parametrize( + "text,expected_name,expected_complete", + [ + ('{"name": "func", "arguments": {"city"', "func", False), + ('{"name": "func", "arguments": {"city": "BJ"}}', "func", None), + ("plain text", None, None), + ('{"nam', None, None), + ('{"name": "deep", "arguments": {"a": {"b": {"c": ', "deep", None), + ('{"name": "batch", "arguments": {"items": [1, 2, ', "batch", None), + ], +) +def test_find_partial_tool_call(text, expected_name, expected_complete): + result = _find_partial_tool_call(text) + if expected_name is None: + assert result is None + else: + assert result is not None + assert result["name"] == expected_name + if expected_complete is not None: + assert result["complete"] is expected_complete def test_feed_plain_text(): @@ -269,7 +215,6 @@ def test_feed_tool_call_args_streaming(): parser = SimpleJsonToolParser() d1 = parser.feed('{"name": "f", "arguments": {"x":') d2 = parser.feed('{"name": "f", "arguments": {"x": "1"}}') - args_deltas = [ d for batch in (d1, d2) @@ -332,17 +277,6 @@ def test_feed_content_after_tool_call_is_not_emitted(): assert parser.has_tool_calls -def _collect_args_deltas(parser): - args_parts = [] - for d in parser.feed(parser._text_buffer): - if "tool_calls" in d: - for tc in d["tool_calls"]: - fn = tc.get("function", {}) - if "arguments" in fn and fn["arguments"]: - args_parts.append(fn["arguments"]) - return args_parts - - def _simulate_streaming(parser, text): all_delta_names = [] all_args_chunks = [] @@ -447,7 +381,6 @@ def test_streaming_args_diff_only_emits_new_bytes(): parser = SimpleJsonToolParser() step1 = parser.feed('{"name": "f", "arguments": {"city": "Bei') step2 = parser.feed('{"name": "f", "arguments": {"city": "Beijing"}}') - all_args = [] for step in (step1, step2): for d in step: @@ -500,31 +433,21 @@ def test_parse_complete_with_content(): def test_parse_complete_multiple_tool_calls(): parser = SimpleJsonToolParser() - body = ( - '{"name": "get_weather", "arguments": {"city": "Beijing"}}' - '{"name": "get_time", "arguments": {"tz": "Asia/Shanghai"}}' - ) + body = '{"name": "get_weather", "arguments": {"city": "Beijing"}}{"name": "get_time", "arguments": {"tz": "Asia/Shanghai"}}' result = parser.parse_complete(body) assert result is not None assert len(result["tool_calls"]) == 2 assert result["tool_calls"][0]["function"]["name"] == "get_weather" assert result["tool_calls"][1]["function"]["name"] == "get_time" - assert "Beijing" in result["tool_calls"][0]["function"]["arguments"] - assert "Asia/Shanghai" in result["tool_calls"][1]["function"]["arguments"] def test_parse_complete_complex_real_world(): parser = SimpleJsonToolParser() body = ( - '{"name": "send_email", ' - '"arguments": {' - '"to": ["a@b.com", "c@d.com"], ' - '"cc": null, ' - '"subject": "Hello World", ' - '"body": "This is a test email.", ' - '"priority": 1, ' - '"attachments": false' - "}}" + '{"name": "send_email", "arguments": {' + '"to": ["a@b.com", "c@d.com"], "cc": null, ' + '"subject": "Hello World", "body": "This is a test email.", ' + '"priority": 1, "attachments": false}}' ) result = parser.parse_complete(body) assert result is not None @@ -539,11 +462,7 @@ def test_parse_complete_complex_real_world(): def test_parse_complete_content_with_multiple_tool_calls(): parser = SimpleJsonToolParser() - body = ( - "I will do two things. " - '{"name": "f1", "arguments": {"a": 1}}' - '{"name": "f2", "arguments": {"b": 2}}' - ) + body = 'I will do two things. {"name": "f1", "arguments": {"a": 1}}{"name": "f2", "arguments": {"b": 2}}' result = parser.parse_complete(body) assert result is not None assert result["content"] == "I will do two things." @@ -588,30 +507,29 @@ def test_feed_then_parse_complete_same_instance(): assert parser.has_tool_calls -def test_pattern_matches_basic(): - assert _TOOL_CALL_HEAD_RE.search('{"name": "f"}') - - -def test_pattern_matches_with_whitespace(): - assert _TOOL_CALL_HEAD_RE.search('{ "name" : "f"}') - - -def test_pattern_no_match_without_name(): - assert _TOOL_CALL_HEAD_RE.search('{"other": 1}') is None - - -def test_pattern_match_mid_text(): - assert _TOOL_CALL_HEAD_RE.search('prefix {"name": "f", "args": {}}') is not None +@pytest.mark.parametrize( + "text,matches", + [ + ('{"name": "f"}', True), + ('{ "name" : "f"}', True), + ('{"other": 1}', False), + ('prefix {"name": "f", "args": {}}', True), + ('{"name": "f"}', True), # match at start + (' {"name": "f"}', True), + ], +) +def test_pattern_regex(text, matches): + result = _TOOL_CALL_HEAD_RE.search(text) + if matches: + assert result is not None + else: + assert result is None def test_pattern_name_at_start(): assert _TOOL_CALL_HEAD_RE.match('{"name": "f"}') -def test_pattern_leading_whitespace(): - assert _TOOL_CALL_HEAD_RE.search(' {"name": "f"}') is not None - - def test_factory_register_and_create(): parser = ToolParserFactory.create("simple_json") assert isinstance(parser, BaseToolParser) @@ -661,7 +579,6 @@ def test_feed_token_ids_do_not_affect_parsing(): text, current_token_ids=[1, 2, 3], delta_token_ids=[3] ) assert len(result_no) == len(result_with) - assert len(result_no) > 0 assert ( result_no[0]["tool_calls"][0]["function"]["name"] == result_with[0]["tool_calls"][0]["function"]["name"] diff --git a/tests/module/test_encoder.py b/tests/module/test_encoder.py index a78e8b3..b66a595 100644 --- a/tests/module/test_encoder.py +++ b/tests/module/test_encoder.py @@ -1,6 +1,13 @@ +import json +import os +import tempfile + +import pytest +import safetensors.torch as st import torch from astrai.config.model_config import EncoderConfig +from astrai.model.automodel import AutoModel from astrai.model.encoder import EmbeddingEncoder TINY_CONFIG = dict( @@ -14,92 +21,56 @@ TINY_CONFIG = dict( norm_eps=1e-5, ) +_device = "cuda" if torch.cuda.is_available() else "cpu" -def test_encoder_forward_mean(): - config = EncoderConfig(**TINY_CONFIG) - device = "cuda" if torch.cuda.is_available() else "cpu" - model = EmbeddingEncoder(config).to(device=device) + +def _make_model(**kwargs): + config = EncoderConfig(**{**TINY_CONFIG, **kwargs}) + return EmbeddingEncoder(config).to(device=_device) + + +@pytest.mark.parametrize("pooling_type", ["mean", "cls", "last"]) +def test_encoder_forward_pooling(pooling_type): + model = _make_model(pooling_type=pooling_type) model.eval() batch_size, seq_len = 2, 8 input_ids = torch.randint( - 0, config.vocab_size, (batch_size, seq_len), device=device + 0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device ) with torch.no_grad(): output = model(input_ids) - assert output.shape == (batch_size, config.dim) - assert not torch.isnan(output).any() - - -def test_encoder_forward_cls(): - config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "cls"}) - device = "cuda" if torch.cuda.is_available() else "cpu" - model = EmbeddingEncoder(config).to(device=device) - model.eval() - - batch_size, seq_len = 2, 8 - input_ids = torch.randint( - 0, config.vocab_size, (batch_size, seq_len), device=device - ) - - with torch.no_grad(): - output = model(input_ids) - - assert output.shape == (batch_size, config.dim) - assert not torch.isnan(output).any() - - -def test_encoder_forward_last(): - config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "last"}) - device = "cuda" if torch.cuda.is_available() else "cpu" - model = EmbeddingEncoder(config).to(device=device) - model.eval() - - batch_size, seq_len = 2, 8 - input_ids = torch.randint( - 0, config.vocab_size, (batch_size, seq_len), device=device - ) - - with torch.no_grad(): - output = model(input_ids) - - assert output.shape == (batch_size, config.dim) + assert output.shape == (batch_size, TINY_CONFIG["dim"]) assert not torch.isnan(output).any() def test_encoder_forward_with_padding(): - config = EncoderConfig(**TINY_CONFIG) - device = "cuda" if torch.cuda.is_available() else "cpu" - model = EmbeddingEncoder(config).to(device=device) + model = _make_model() model.eval() batch_size, seq_len = 2, 8 input_ids = torch.randint( - 0, config.vocab_size, (batch_size, seq_len), device=device + 0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device ) - input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=device) + input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=_device) input_mask[:, 4:] = False with torch.no_grad(): output = model(input_ids, input_mask=input_mask) - assert output.shape == (batch_size, config.dim) + assert output.shape == (batch_size, TINY_CONFIG["dim"]) assert not torch.isnan(output).any() def test_encoder_normalize(): - config = EncoderConfig( - **{**TINY_CONFIG, "pooling_type": "mean", "normalize_embeddings": True} - ) - device = "cuda" if torch.cuda.is_available() else "cpu" - model = EmbeddingEncoder(config).to(device=device) + model = _make_model(pooling_type="mean", normalize_embeddings=True) model.eval() batch_size, seq_len = 2, 8 input_ids = torch.randint( - 0, config.vocab_size, (batch_size, seq_len), device=device + 0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device ) with torch.no_grad(): @@ -110,24 +81,19 @@ def test_encoder_normalize(): def test_encoder_register(): - from astrai.model.automodel import AutoModel - assert AutoModel.is_registered("embedding") cls = AutoModel.get_component_class("embedding") assert cls is EmbeddingEncoder def test_encoder_from_transformer_checkpoint(): - config = EncoderConfig(**TINY_CONFIG) - device = "cuda" if torch.cuda.is_available() else "cpu" - model = EmbeddingEncoder(config).to(device=device) - + model = _make_model() state_dict = model.state_dict() state_dict["lm_head.weight"] = torch.randn( - config.vocab_size, config.dim, device=device + TINY_CONFIG["vocab_size"], TINY_CONFIG["dim"], device=_device ) - new_model = EmbeddingEncoder(config).to(device=device) + new_model = _make_model() new_model.load_state_dict(state_dict, strict=True) for key in model.state_dict(): @@ -135,12 +101,6 @@ def test_encoder_from_transformer_checkpoint(): def test_encoder_save_load(): - import json - import os - import tempfile - - import safetensors.torch as st - test_dir = tempfile.mkdtemp(prefix="encoder_test_") config_path = os.path.join(test_dir, "config.json") weights_path = os.path.join(test_dir, "model.safetensors")