refactor : 压缩测试代码，消除重复
- fixture 替代重复实例化和 tokenizer 落盘 - parametrize 合并同构测试 - helper 消除 save_h5 + DatasetFactory.load 样板 - 净减 272 行
This commit is contained in:
parent
39985840c7
commit
25d4ea3f91
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -13,67 +13,28 @@ 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 complete is expected_complete
|
||||
if check_end_eq_len:
|
||||
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
|
||||
|
||||
|
||||
def test_find_single_tool_call():
|
||||
text = '{"name": "get_weather", "arguments": {"city": "Beijing"}}'
|
||||
results = _find_tool_calls(text)
|
||||
|
|
@ -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"')
|
||||
@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"] == "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"
|
||||
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"]
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Reference in New Issue