| import json
|
|
|
| import numpy as np
|
| import pytest
|
| import torch
|
| from safetensors.torch import save_file
|
|
|
| from laya_onnx.inputs import collate_items
|
| from laya_onnx.torch_model import DecisionModel
|
|
|
|
|
| @pytest.fixture
|
| def checkpoint(tmp_path):
|
| torch.manual_seed(2026)
|
| cfg = {
|
| "model_type": "modernbert",
|
| "hidden_size": 32,
|
| "intermediate_size": 48,
|
| "vocab_size": 100,
|
| "num_hidden_layers": 3,
|
| "num_attention_heads": 4,
|
| "local_attention": 8,
|
| "global_attn_every_n_layers": 3,
|
| }
|
| agent = {
|
| "head_layers": 2,
|
| "act_costs": {"escalate": 0.5},
|
| "max_len": 64,
|
| "head_max_len": 32,
|
| "encoder": "tiny-test",
|
| "temperature": [1.0, 1.0, 1.0],
|
| }
|
| model = DecisionModel(cfg, agent, 64).eval()
|
| for name, p in model.named_parameters():
|
| if "in_proj" in name:
|
| torch.nn.init.normal_(p, std=0.03)
|
| (tmp_path / "encoder").mkdir()
|
| (tmp_path / "encoder/config.json").write_text(json.dumps(cfg))
|
| (tmp_path / "rl_agent_config.json").write_text(json.dumps(agent))
|
| (tmp_path / "tokenizer").mkdir()
|
| from tokenizers import Tokenizer
|
| from tokenizers.models import WordLevel
|
|
|
| tok = Tokenizer(
|
| WordLevel(
|
| {"[UNK]": 0, "[CLS]": 1, "[SEP]": 2, "[MASK]": 3, "[PAD]": 4},
|
| unk_token="[UNK]",
|
| )
|
| )
|
| tok.save(str(tmp_path / "tokenizer/tokenizer.json"))
|
| (tmp_path / "tokenizer/tokenizer_config.json").write_text(
|
| json.dumps(
|
| {
|
| "cls_token": "[CLS]",
|
| "sep_token": "[SEP]",
|
| "mask_token": "[MASK]",
|
| "pad_token": "[PAD]",
|
| }
|
| )
|
| )
|
| save_file(model.state_dict(), str(tmp_path / "model.safetensors"))
|
| return tmp_path, model
|
|
|
|
|
| def batch(length=31, batch_size=3):
|
| items = [
|
| {"ids": list(range(length)), "markers": [3, 8, 10][: i + 1], "qtype": i}
|
| for i in range(batch_size)
|
| ]
|
| shape = {"batch_size": 3, "max_length": 64, "min_length": 16, "max_options": 4}
|
| return {k: torch.from_numpy(v) for k, v in collate_items(items, 4, shape=shape).items()}
|
|
|
|
|
| def test_onnx_export_parity_and_dynamic_lengths(checkpoint):
|
| from laya_onnx import Agent
|
| from laya_onnx.convert import convert
|
|
|
| source, model = checkpoint
|
| output = source / "export"
|
| convert(source, output, batch_size=3, max_options=4)
|
| agent = Agent(output)
|
| for length in (17, 47, 64):
|
| inputs = batch(length)
|
| with torch.inference_mode():
|
| expected = model(**inputs)
|
| actual = agent.forward({k: v.numpy() for k, v in inputs.items()})
|
| for a, e in zip(actual, expected):
|
| np.testing.assert_allclose(a, e.numpy(), atol=1e-4, rtol=1e-4)
|
| result = agent.predict(
|
| "sample",
|
| {"q": {"type": "choice", "instructions": "pick", "criteria": ["a", "b"]}},
|
| )
|
| assert result["usage"]["output_tokens"] == 0
|
| assert set(result["answers"]["q"]["probabilities"]) == {"a", "b"}
|
| with pytest.raises(FileExistsError):
|
| convert(source, output)
|
|
|
