File size: 5,770 Bytes
5f70891 | 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 115 116 | """Research recognizer sharing one English encoder across two explicit branches.
The character branch applies saved low-rank corrections through functional_call
and supplies precomputed encoder outputs to a decoder-only VED container. No
second learned encoder is registered or saved. Functional parameter overrides
are serial CPU research operations, not a concurrent/native serving interface.
"""
import copy
import hashlib
import json
from pathlib import Path
import torch
from torch import nn
from safetensors.torch import load_file
from transformers import DeiTModel, GenerationConfig, VisionEncoderDecoderConfig, VisionEncoderDecoderModel
def sha(path):
value = hashlib.sha256()
with Path(path).open('rb') as stream:
for block in iter(lambda: stream.read(1024 * 1024), b''): value.update(block)
return value.hexdigest()
def eager(config):
config._attn_implementation = 'eager'
config.encoder._attn_implementation = 'eager'
config.decoder._attn_implementation = 'eager'
return config
def load_english(folder):
config = eager(VisionEncoderDecoderConfig.from_pretrained(folder, local_files_only=True))
model, report = VisionEncoderDecoderModel.from_pretrained(folder, config=config,
encoder=DeiTModel(config.encoder, add_pooling_layer=False), local_files_only=True,
dtype=torch.float32, attn_implementation='eager', output_loading_info=True)
assert not any(report.get(k) for k in ['missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs'])
assert sum(p.numel() for p in model.parameters()) == 61448832
return model.eval().requires_grad_(False)
class PrecomputedOnlyEncoder(nn.Module):
main_input_name = 'pixel_values'
def __init__(self, config):
super().__init__(); self.config = copy.deepcopy(config)
def get_output_embeddings(self):
return None
def forward(self, pixel_values=None, **kwargs):
raise RuntimeError('Character decoding requires the shared encoder outputs')
class Correction(nn.Module):
def __init__(self, spec):
super().__init__(); self.spec = spec
if spec['kind'] == 'factor':
self.left = nn.Parameter(torch.empty(spec['matrix_shape'][0], spec['rank']))
self.right = nn.Parameter(torch.empty(spec['rank'], spec['matrix_shape'][1]))
else:
assert spec['kind'] == 'replacement'
self.value = nn.Parameter(torch.empty(spec['shape']))
def apply(self, base):
if self.spec['kind'] == 'replacement': return self.value
return base + (self.left @ self.right).reshape(self.spec['shape'])
class SharedCharacterRecognizer(nn.Module):
def __init__(self, original, character_decoder, character_projection, character_config,
english_generation, character_generation, specifications):
super().__init__()
self.original = original
cfg = eager(copy.deepcopy(character_config))
self.character = VisionEncoderDecoderModel(config=cfg,
encoder=PrecomputedOnlyEncoder(cfg.encoder), decoder=character_decoder)
if character_projection is not None: self.character.enc_to_dec_proj = character_projection
self.original.generation_config = copy.deepcopy(english_generation)
self.character.generation_config = copy.deepcopy(character_generation)
self.specifications = specifications
self.corrections = nn.ModuleDict({spec['key']: Correction(spec) for spec in specifications})
def generate_pixels(self, pixels, branch, generation):
if self.training: raise ValueError('Explicit eval mode is required')
if branch == 'english': return self.original.generate(pixel_values=pixels, **generation)
if branch != 'character': raise ValueError('Require an explicit English or character branch')
base = self.original.encoder.state_dict()
overrides = {spec['name']: self.corrections[spec['key']].apply(base[spec['name']])
for spec in self.specifications}
output = torch.func.functional_call(self.original.encoder, overrides, (),
dict(pixel_values=pixels, return_dict=True), strict=False)
return self.character.generate(encoder_outputs=output, **generation)
@classmethod
def from_bundle(cls, folder):
folder = Path(folder); manifest = json.loads((folder / 'bundle.json').read_text())
for name, expected in manifest['files_sha256'].items(): assert sha(folder / name) == expected
configs = json.loads((folder / 'model-configs.json').read_text())
english = eager(VisionEncoderDecoderConfig.from_dict(configs['english']))
character = eager(VisionEncoderDecoderConfig.from_dict(configs['character']))
original = VisionEncoderDecoderModel(config=english,
encoder=DeiTModel(english.encoder, add_pooling_layer=False))
private = VisionEncoderDecoderModel(config=character,
encoder=PrecomputedOnlyEncoder(character.encoder))
model = cls(original, private.decoder, getattr(private, 'enc_to_dec_proj', None), character,
GenerationConfig.from_pretrained(folder, config_file_name='english-generation.json', local_files_only=True),
GenerationConfig.from_pretrained(folder, config_file_name='character-generation.json', local_files_only=True),
manifest['corrections'])
result = model.load_state_dict(load_file(str(folder / 'weights.safetensors')), strict=True)
assert not result.missing_keys and not result.unexpected_keys
assert sum(p.numel() for p in model.parameters()) == manifest['parameters'] == 74842752
assert len(model.character.encoder.state_dict()) == 0
return model.eval().requires_grad_(False)
|