ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: deduplicate and restructure test suite

Browse Source 
- extract preprocessing config factories into tests/data/factories.py
- keep conftest.py fixtures-only; stop importing builders from it
- promote temp_dir fixture to root conftest for cross-directory reuse
- unify duplicate BPE tokenizer builders into build_test_tokenizer
- merge grpo/dpo online e2e tests into one parametrized integration test
- extract engine mock factory and shared model batch builders
- drop local tempfile usage in favor of shared fixtures

No behavior change: 519 tests pass.
This commit is contained in:
ViperEkura
2026-08-20 01:53:28 +08:00
parent 53a7149577
commit 84753d3e08
13 changed files with 312 additions and 454 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/data/conftest.py
+44 -169
View File
@@ -1,21 +1,15 @@
import json
import os
import tempfile
import pytest
from tokenizers import Tokenizer, models, pre_tokenizers, trainers
from astrai.config.preprocess_config import (
InputConfig,
PipelineConfig,
ProcessingConfig,
)
from astrai.preprocessing.builder import (
MultiOutputMaskBuilder,
SectionedMaskBuilder,
SingleOutputMaskBuilder,
)
from astrai.tokenize import AutoTokenizer
from tests.data.factories import make_grpo_config
from tests.helpers import build_test_tokenizer
_SPECIAL_TOKENS_CONFIG = {
"bos_token": "<|begin_of_sentence|>",
@@ -41,55 +35,42 @@ _CHAT_TEMPLATE = (
"{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
)
_CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}]
_INSTRUCTION_SECTIONS = [
{"field": "prompt", "action": "mask", "add_special_tokens": True},
{"field": "response", "action": "train"},
_CHAT_TOKENIZER_DATA = [
"hello world",
"Hi there!",
"You are helpful.",
"What is 2+2?",
"Tell me a story about dragons and knights.",
"Sure, here is a tale.",
"Translate to French: Hello",
"Bonjour",
"Artificial Intelligence is a field of computer science.",
"system",
"user",
"assistant",
"<|im_start|>",
"<|im_end|>",
*[chr(i) for i in range(32, 127)],
]
_TEXT_SECTIONS = [{"field": "text", "action": "train"}]
_GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}]
_CHAT_TOKENIZER_MAP = {
"bos_token": "<|begin_of_sentence|>",
"eos_token": "<|end_of_sentence|>",
"pad_token": "<|_pad_|>",
"unk_token": "<|_unk_|>",
}
def _build_chat_tokenizer():
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tr = trainers.BpeTrainer(
return build_test_tokenizer(
vocab_size=512,
min_frequency=1,
special_tokens=_SPECIAL_TOKENS,
special_token_map=_CHAT_TOKENIZER_MAP,
add_prefix_space=False,
train_data=_CHAT_TOKENIZER_DATA,
chat_template=_CHAT_TEMPLATE,
)
train_data = [
"hello world",
"Hi there!",
"You are helpful.",
"What is 2+2?",
"Tell me a story about dragons and knights.",
"Sure, here is a tale.",
"Translate to French: Hello",
"Bonjour",
"Artificial Intelligence is a field of computer science.",
"system",
"user",
"assistant",
"<|im_start|>",
"<|im_end|>",
*[chr(i) for i in range(32, 127)],
]
tok.train_from_iterator(train_data, tr)
auto_tok = AutoTokenizer()
auto_tok._tokenizer = tok
auto_tok._special_token_map = {
"bos_token": "<|begin_of_sentence|>",
"eos_token": "<|end_of_sentence|>",
"pad_token": "<|_pad_|>",
"unk_token": "<|_unk_|>",
}
auto_tok.set_chat_template(_CHAT_TEMPLATE)
return auto_tok
@pytest.fixture(scope="session")
@@ -97,116 +78,11 @@ def chat_tokenizer():
return _build_chat_tokenizer()
@pytest.fixture
def temp_dir():
d = tempfile.mkdtemp()
yield d
import shutil
shutil.rmtree(d, ignore_errors=True)
def make_chat_config():
return PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
mask={"system": "mask", "user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def make_instruction_config():
return PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def make_text_config():
return PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
preprocessing=ProcessingConfig(
max_seq_len=2048, min_chars=1, max_chars=2_000_000
),
)
def make_dpo_chat_config():
return PipelineConfig(
input=InputConfig(
sources={
"chosen": {
"sections": [
{"field": "chosen", "action": "$role", "template": True}
]
},
"rejected": {
"sections": [
{"field": "rejected", "action": "$role", "template": True}
]
},
}
),
mask={"user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def make_grpo_config():
return PipelineConfig(
input=InputConfig(
sources={
"prompts": {
"sections": [
{"field": "prompt", "action": "mask", "template": True}
]
},
"responses": {
"sections": _GRPO_RESPONSE_SECTIONS,
"list_field": True,
"mask_key": "masks",
},
"rewards": {
"sections": [{"field": "rewards", "action": "value"}],
},
}
),
mask={"user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def make_grpo_no_template_config():
return PipelineConfig(
input=InputConfig(
sources={
"prompts": {
"sections": [
{
"field": "prompt",
"action": "mask",
"add_special_tokens": True,
}
]
},
"responses": {
"sections": _GRPO_RESPONSE_SECTIONS,
"list_field": True,
"mask_key": "masks",
},
"rewards": {
"sections": [{"field": "rewards", "action": "value"}],
},
}
),
mask={"user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def _write_tokenizer_dir(dir_path, tokenizer, tokenizer_config):
"""Persist a tokenizer plus ``tokenizer_config.json`` into *dir_path*."""
tokenizer._tokenizer.save(os.path.join(dir_path, "tokenizer.json"))
with open(os.path.join(dir_path, "tokenizer_config.json"), "w") as f:
json.dump(tokenizer_config, f)
@pytest.fixture
@@ -228,11 +104,11 @@ def multi_builder():
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
)
_write_tokenizer_dir(
d,
test_tokenizer,
{"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}},
)
return d
@@ -240,10 +116,9 @@ def tokenizer_dir(temp_dir, test_tokenizer):
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,
)
_write_tokenizer_dir(
d,
chat_tokenizer,
{"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE},
)
return d
tests/data/factories.py
+86
View File
@@ -0,0 +1,86 @@
"""Test data builders for preprocessing and dataset scenarios."""
from astrai.config.preprocess_config import (
InputConfig,
PipelineConfig,
ProcessingConfig,
)
CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}]
INSTRUCTION_SECTIONS = [
{"field": "prompt", "action": "mask", "add_special_tokens": True},
{"field": "response", "action": "train"},
]
TEXT_SECTIONS = [{"field": "text", "action": "train"}]
GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}]
def make_pipeline_config(sections, *, mask=None, preprocessing=None, sources=None):
"""Build a pipeline config with the common test defaults."""
return PipelineConfig(
input=InputConfig(sections=sections, sources=sources),
mask={} if mask is None else mask,
mask_default="mask",
preprocessing=preprocessing or ProcessingConfig(max_seq_len=2048),
)
def make_chat_config():
return make_pipeline_config(
CHAT_SECTIONS,
mask={"system": "mask", "user": "mask", "assistant": "train"},
)
def make_instruction_config():
return make_pipeline_config(
INSTRUCTION_SECTIONS,
mask={"prompt": "mask", "response": "train"},
)
def make_text_config():
return make_pipeline_config(
TEXT_SECTIONS,
preprocessing=ProcessingConfig(
max_seq_len=2048, min_chars=1, max_chars=2_000_000
),
)
def make_dpo_chat_config():
sources = {
name: {"sections": [{"field": name, "action": "$role", "template": True}]}
for name in ("chosen", "rejected")
}
return make_pipeline_config(
None,
mask={"user": "mask", "assistant": "train"},
sources=sources,
)
def make_grpo_config(*, template=True):
prompt_section = {"field": "prompt", "action": "mask"}
if template:
prompt_section["template"] = True
else:
prompt_section["add_special_tokens"] = True
sources = {
"prompts": {"sections": [prompt_section]},
"responses": {
"sections": GRPO_RESPONSE_SECTIONS,
"list_field": True,
"mask_key": "masks",
},
"rewards": {"sections": [{"field": "rewards", "action": "value"}]},
}
return make_pipeline_config(
None,
mask={"user": "mask", "assistant": "train"},
sources=sources,
)
def make_grpo_no_template_config():
return make_grpo_config(template=False)
tests/data/test_dataset.py
+2 -2
View File
@@ -24,7 +24,7 @@ from astrai.serialization import (
load_bin,
save_bin,
)
from tests.data.conftest import make_grpo_no_template_config
from tests.data.factories import make_grpo_config
def _rand_seq(length, vocab=1000):
@@ -797,7 +797,7 @@ def test_grpo_builder_preserves_response_boundaries(base_test_env):
_save_test_tokenizer(base_test_env["test_dir"], tokenizer)
builder = SectionedMaskBuilder()
config = make_grpo_no_template_config()
config = make_grpo_config(template=False)
config.preprocessing.max_seq_len = 128
item = {
tests/data/test_preprocess_builder.py
+13 -13
View File
@@ -12,10 +12,10 @@ from astrai.preprocessing.builder import (
SectionedMaskBuilder,
SingleOutputMaskBuilder,
)
from tests.data.conftest import (
_CHAT_SECTIONS,
_INSTRUCTION_SECTIONS,
_TEXT_SECTIONS,
from tests.data.factories import (
CHAT_SECTIONS,
INSTRUCTION_SECTIONS,
TEXT_SECTIONS,
make_chat_config,
make_dpo_chat_config,
make_grpo_config,
@@ -101,7 +101,7 @@ def test_chat_uniform_masking(
mask_rules, mask_default, expect_nonzero, chat_tokenizer, builder
):
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask=mask_rules,
mask_default=mask_default,
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -128,7 +128,7 @@ def test_chat_empty_messages(chat_tokenizer, builder):
def test_chat_domain_extraction(chat_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask={"assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -147,7 +147,7 @@ def test_chat_domain_extraction(chat_tokenizer, builder):
def test_chat_truncation(chat_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask={"assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=10),
@@ -237,7 +237,7 @@ def test_text_empty(test_tokenizer, builder):
def test_text_too_short(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(min_chars=100),
)
assert builder.build({"text": "short"}, config, test_tokenizer) is None
@@ -245,7 +245,7 @@ def test_text_too_short(test_tokenizer, builder):
def test_text_truncation(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1),
)
item = {"text": "This is a very long text that should be truncated"}
@@ -255,7 +255,7 @@ def test_text_truncation(test_tokenizer, builder):
def test_sectioned_chat(chat_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask={"system": "mask", "user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -275,7 +275,7 @@ def test_sectioned_chat(chat_tokenizer, builder):
def test_sectioned_instruction(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
input=InputConfig(sections=INSTRUCTION_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0),
)
item = {"prompt": "Q: Why?", "response": "A: Because."}
@@ -288,7 +288,7 @@ def test_sectioned_instruction(test_tokenizer, builder):
def test_sectioned_text(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1),
)
item = {"text": "Hello world, this is a test."}
@@ -299,7 +299,7 @@ def test_sectioned_text(test_tokenizer, builder):
def test_sectioned_text_too_short(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100),
)
assert builder.build({"text": "short"}, config, test_tokenizer) is None
tests/data/test_preprocess_config.py
+7 -7
View File
@@ -4,9 +4,9 @@ from astrai.config.preprocess_config import (
InputConfig,
PipelineConfig,
)
from tests.data.conftest import (
_INSTRUCTION_SECTIONS,
_TEXT_SECTIONS,
from tests.data.factories import (
INSTRUCTION_SECTIONS,
TEXT_SECTIONS,
make_dpo_chat_config,
)
@@ -43,26 +43,26 @@ def test_from_dict_flat():
def test_to_dict_roundtrip():
config = PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
input=InputConfig(sections=INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
)
d = config.to_dict()
config2 = PipelineConfig.from_dict(d)
assert config2.input.sections == _INSTRUCTION_SECTIONS
assert config2.input.sections == INSTRUCTION_SECTIONS
assert config2.mask == {"prompt": "mask", "response": "train"}
def test_to_file_from_file(temp_dir):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
mask={"text": "train"},
mask_default="mask",
)
path = os.path.join(temp_dir, "config.json")
config.to_file(path)
loaded = PipelineConfig.from_file(path)
assert loaded.input.sections == _TEXT_SECTIONS
assert loaded.input.sections == TEXT_SECTIONS
assert loaded.mask == {"text": "train"}
tests/data/test_preprocess_pipeline.py
+8 -8
View File
@@ -9,10 +9,10 @@ from astrai.config.preprocess_config import (
)
from astrai.preprocessing.packing import PackingStrategyFactory
from astrai.preprocessing.pipeline import Pipeline, filter_by_length
from tests.data.conftest import (
_CHAT_SECTIONS,
_INSTRUCTION_SECTIONS,
_TEXT_SECTIONS,
from tests.data.factories import (
CHAT_SECTIONS,
INSTRUCTION_SECTIONS,
TEXT_SECTIONS,
make_dpo_chat_config,
make_grpo_no_template_config,
)
@@ -54,7 +54,7 @@ def test_full_chat_pipeline(temp_dir, chat_tokenizer_dir):
)
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask={"system": "mask", "user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -97,7 +97,7 @@ def test_full_text_pipeline(temp_dir, tokenizer_dir):
)
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=10),
output=OutputConfig(storage_format="bin"),
)
@@ -138,7 +138,7 @@ def test_full_instruction_pipeline(temp_dir, tokenizer_dir):
)
config = PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
input=InputConfig(sections=INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -164,7 +164,7 @@ def test_dtype_override(temp_dir, tokenizer_dir):
f.write(json.dumps({"prompt": "Q", "response": "A"}) + "\n")
config = PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
input=InputConfig(sections=INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),