File size: 6,083 Bytes
ae36cfb | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """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()
|