| """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() |
|
|
