| """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) |
|
|
