File size: 1,860 Bytes
f36843f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | """Check the published SDK router and adapter without downloading weights."""
import unittest
from unittest.mock import patch
import torch
from transformers import Qwen2Config, Qwen2ForCausalLM
from core.engine import run_parallel_generation
from core.schema import StructuredSchema
class Tokenizer:
pad_token_id = 0
def encode(self, text, return_tensors=None, **kwargs):
ids = [1 + ord(c) % 95 for c in text]
return torch.tensor([ids]) if return_tensors == 'pt' else ids
class SDKTest(unittest.TestCase):
def test_tree_is_default_and_preserves_response(self):
torch.set_num_threads(2)
torch.manual_seed(42)
config = Qwen2Config(vocab_size=97, hidden_size=64, intermediate_size=128,
num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2)
config._attn_implementation = 'sdpa'
model = Qwen2ForCausalLM(config).eval()
schema = StructuredSchema({'risk': {'type':'boolean', 'description':'Risk present'},
'level': {'type':'enum', 'description':'Risk level', 'choices':['LOW','HIGH']}})
self.assertEqual(run_parallel_generation.__module__, 'core.engine_tree')
with patch('core.engine_tree.get_torch_engine', return_value=(model, Tokenizer(), 'cpu')):
result = run_parallel_generation('A suspicious request.', schema)
self.assertEqual(result['mode'], 'parallel_constrained_tree')
self.assertEqual(set(result['parsed_json']), {'risk', 'level'})
self.assertIsInstance(result['parsed_json']['risk']['value'], bool)
self.assertIn(result['parsed_json']['level']['value'], ['LOW', 'HIGH'])
self.assertFalse(result['has_calibrated_probabilities'])
self.assertEqual(result['num_fields'], 2)
if __name__ == '__main__':
unittest.main()
|