| """Check training-only head initialization without weakening predictor loading.""" |
| import importlib.util |
| from pathlib import Path |
| import tempfile |
| import unittest |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| spec = importlib.util.spec_from_file_location('regression_loader', ROOT / 'validate_release.py') |
| loader = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(loader) |
|
|
|
|
| class RegressionBaseTests(unittest.TestCase): |
| def test_shared_default(self): |
| self.assertEqual(loader.REGRESSION_BASE, 'RNASeek-regression-pretrained') |
| self.assertTrue(loader.local_path(loader.REGRESSION_BASE).is_dir()) |
| self.assertTrue(loader.local_path('RNASeek-lm-pretrained').is_dir()) |
| for task in loader.DEFAULTS.values(): |
| self.assertNotEqual(task['checkpoint'], loader.REGRESSION_BASE) |
|
|
| def test_pretrained_and_predictor_loading(self): |
| import torch |
| from safetensors.torch import save_file |
| from tokenizers import Tokenizer, models |
| from transformers import Qwen2Config, Qwen2ForCausalLM, PreTrainedTokenizerFast |
| with tempfile.TemporaryDirectory() as temporary: |
| base = Path(temporary) / 'base' |
| config = Qwen2Config(vocab_size=8, hidden_size=8, intermediate_size=16, |
| num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=1, |
| max_position_embeddings=32, pad_token_id=0, bos_token_id=1, eos_token_id=2) |
| original = Qwen2ForCausalLM(config) |
| original.save_pretrained(base) |
| tokenizer = PreTrainedTokenizerFast(tokenizer_object=Tokenizer(models.WordLevel( |
| {'<pad>': 0, '<s>': 1, '</s>': 2, '<unk>': 3, 'A': 4}, unk_token='<unk>')), |
| pad_token='<pad>', bos_token='<s>', eos_token='</s>', unk_token='<unk>') |
| tokenizer.save_pretrained(base) |
| with self.assertRaisesRegex(ValueError, 'no trained regression head'): |
| loader.load_regression_checkpoint(base) |
| backbone, head, _ = loader.load_regression_checkpoint(base, initialize_head=True) |
| for key, value in original.model.state_dict().items(): |
| torch.testing.assert_close(backbone.state_dict()[key], value, rtol=0, atol=0) |
| head.net(torch.randn(2, 8)).sum().backward() |
| self.assertIsNotNone(head.net[-1].weight.grad) |
| predictor = Path(temporary) / 'predictor' |
| predictor.mkdir() |
| state = {'backbone.' + k: v for k, v in backbone.state_dict().items()} |
| state.update({'regression_head.' + k: v for k, v in head.state_dict().items()}) |
| save_file(state, str(predictor / 'model.safetensors')) |
| config.save_pretrained(predictor) |
| tokenizer.save_pretrained(predictor) |
| for training in (False, True): |
| _, retained_head, _ = loader.load_regression_checkpoint(predictor, initialize_head=training) |
| for key, value in head.state_dict().items(): |
| torch.testing.assert_close(retained_head.state_dict()[key], value, rtol=0, atol=0) |
|
|
