File size: 3,098 Bytes
7155827 db88907 7155827 db88907 7155827 | 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 | """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)
|