| """CPU checks for leakage guards, whole-project token masks, and streamed LoRA math.""" |
| import importlib.util |
| import json |
| import tempfile |
| import unittest |
| from pathlib import Path |
| from unittest import mock |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
| def module(name): |
| spec = importlib.util.spec_from_file_location(name, ROOT / 'scripts' / (name + '.py')) |
| obj = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(obj) |
| return obj |
|
|
| train = module('train') |
| export = module('export_model') |
|
|
| class FakeTokenizer: |
| def apply_chat_template(self, messages, tokenize, add_generation_prompt): |
| prefix = [1, 2, 3] |
| return prefix if add_generation_prompt else prefix + [4, 5, 6, 7] |
|
|
| class TrainingTests(unittest.TestCase): |
| def row(self, instruction='one', marker='one'): |
| return {'instruction': instruction, 'output': json.dumps({'files': {'project.godot': marker, 'main.gd': 'extends Node'}})} |
|
|
| def test_assistant_mask_and_whole_project_rejection(self): |
| rows, rejected = train.tokenize_rows([self.row()], FakeTokenizer(), 7, 'train') |
| self.assertEqual(rows[0]['labels'], [-100, -100, -100, 4, 5, 6, 7]) |
| self.assertEqual(rejected, []) |
| with self.assertRaisesRegex(ValueError, 'No complete'): |
| train.tokenize_rows([self.row()], FakeTokenizer(), 6, 'train') |
|
|
| def test_blind_and_output_leakage(self): |
| blind = [{'instruction': f'blind {i}'} for i in range(500)] |
| train.check_splits([self.row('training')], [self.row('development', 'two')], blind) |
| with self.assertRaisesRegex(ValueError, 'contamination'): |
| train.check_splits([self.row('blind 1')], [self.row('development', 'two')], blind) |
| with self.assertRaisesRegex(ValueError, 'Identical project'): |
| train.check_splits([self.row('training')], [self.row('development')], blind) |
| with self.assertRaisesRegex(ValueError, '500'): |
| train.check_splits([self.row()], [self.row('dev', 'other')], blind[:-1]) |
|
|
| def test_only_intact_complete_checkpoint_resumes(self): |
| with tempfile.TemporaryDirectory() as d: |
| root = Path(d) |
| old = root / 'checkpoint-1'; old.mkdir() |
| new = root / 'checkpoint-2'; new.mkdir() |
| (old / 'adapter').write_text('complete') |
| train.atomic_json(old / 'COMPLETE.json', {'sha256': {'adapter': train.digest(old / 'adapter')}}) |
| self.assertEqual(train.latest_complete_checkpoint(root), old) |
| (new / 'adapter').write_text('complete') |
| train.atomic_json(new / 'COMPLETE.json', {'sha256': {'adapter': train.digest(new / 'adapter')}}) |
| (new / 'adapter').write_text('interrupted') |
| self.assertEqual(train.latest_complete_checkpoint(root), old) |
|
|
| class ExportTests(unittest.TestCase): |
| def test_converter_installs_only_into_overlay_and_preserves_torch(self): |
| versions = {'python': '3.12.12', 'packages': {'torch': '2.11.0+cu130'}} |
| with tempfile.TemporaryDirectory() as d: |
| workspace = Path(d) |
| with mock.patch('importlib.metadata.version', return_value='2.11.0+cu130'), \ |
| mock.patch.object(export, 'run') as run, \ |
| mock.patch.object(export.subprocess, 'check_output', return_value=json.dumps(versions)): |
| python, env, actual = export.converter_environment(workspace / 'llama.cpp', workspace) |
| commands = [call.args[0] for call in run.call_args_list] |
| self.assertIn('--system-site-packages', commands[0]) |
| self.assertIn('--without-pip', commands[0]) |
| self.assertEqual(commands[1][:2], [python, '-c']) |
| self.assertIn('sys.prefix != sys.base_prefix', commands[1][2]) |
| self.assertIn('sysconfig.get_path', commands[1][2]) |
| self.assertIn("find_spec('pip')", commands[1][2]) |
| self.assertEqual(commands[2][:4], [python, '-m', 'pip', 'install']) |
| self.assertEqual(python, workspace / 'converter-venv/bin/python') |
| self.assertEqual((workspace / 'converter-constraints.txt').read_text(), 'torch===2.11.0+cu130\n') |
| self.assertEqual(env['PIP_REQUIRE_VIRTUALENV'], 'true') |
| self.assertEqual(env['OMP_NUM_THREADS'], '1') |
| self.assertEqual(actual, versions) |
| with mock.patch('importlib.metadata.version', return_value='2.10.0'), \ |
| mock.patch.object(export, 'run') as run: |
| with self.assertRaisesRegex(RuntimeError, 'requires installed torch'): |
| export.converter_environment(Path('/unused'), Path('/unused')) |
| run.assert_not_called() |
|
|
| def test_row_chunked_merge_matches_dense_math(self): |
| import torch |
| torch.manual_seed(3) |
| w = torch.randn(513, 257, dtype=torch.bfloat16) |
| a = torch.randn(32, 257) * .01 |
| b = torch.randn(513, 32) * .01 |
| merged = export.merge_tensor(w, {'A': a, 'B': b}, 2, row_chunk=41) |
| expected = (w.float() + (b @ a) * 2).half() |
| torch.testing.assert_close(merged, expected, rtol=.001, atol=.001) |
| self.assertEqual(merged.dtype, torch.float16) |
|
|
| def test_adapter_names_and_nonfinite_are_rejected(self): |
| import torch |
| from safetensors.torch import save_file |
| with tempfile.TemporaryDirectory() as d: |
| path = Path(d) / 'adapter.safetensors' |
| save_file({'base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight': torch.zeros(32, 3), |
| 'base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight': torch.zeros(4, 32)}, path) |
| self.assertIn('model.layers.0.self_attn.q_proj.weight', export.lora_pairs(path)) |
| save_file({'unknown': torch.zeros(2)}, path) |
| with self.assertRaisesRegex(ValueError, 'Unsupported'): |
| export.lora_pairs(path) |
| with self.assertRaisesRegex(ValueError, 'Nonfinite'): |
| export.merge_tensor(torch.zeros(3, 3), {'A': torch.ones(1, 3), 'B': torch.full((3, 1), float('inf'))}, 2) |
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
