| """Serial CPU Paddle research runtime for the three shared recognizer graphs. |
| |
| Consumes only native graph parameters, codec/tokenizer assets and the compiled |
| encoder operator library. No Torch weights, source logits, source caches or |
| forced token histories enter free generation. Explicit branch selection only. |
| """ |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import paddle |
| from paddle.utils.cpp_extension import load_op_meta_info_and_register_op |
| import sentencepiece as spm |
| from PIL import Image |
|
|
| from trocr_character_codec import CharacterCodec |
| from trocr_numpy_beam import TrOCRBeamState |
|
|
|
|
| _LIBRARY_HASH=None |
|
|
|
|
| def preprocess_rgb(image): |
| image=image.convert('RGB').resize((384,384),resample=Image.Resampling.BICUBIC) |
| pixels=(np.asarray(image).astype(np.float64)*(1./255.)).astype(np.float32) |
| return np.ascontiguousarray(((pixels-np.float32(.5))/np.float32(.5)).transpose(2,0,1)[None]) |
|
|
|
|
| def read(path):return json.loads(Path(path).read_text()) |
| def sha(path): |
| value=hashlib.sha256() |
| with Path(path).open('rb') as stream: |
| for part in iter(lambda:stream.read(1024*1024),b''):value.update(part) |
| return value.hexdigest() |
|
|
|
|
| class SharedCharacterNative: |
| def __init__(self,candidate,encoder,decoders,allow_unqualified=False,library_path=None,allow_control=False): |
| if not allow_unqualified:raise ValueError('Source research model; full native/PaddleX acceptance is pending') |
| if paddle.device.get_device()!='cpu':raise ValueError('Only the validated CPU operator profile is supported') |
| candidate,encoder,decoders=map(Path,[candidate,encoder,decoders]) |
| manifest=read(candidate/'bundle.json');enc=read(encoder/'report.json');enc_run=read(encoder/'run.json');dec=read(decoders/'report.json') |
| digest=sha(candidate/'bundle.json') |
| assert digest==enc_run['candidate_manifest_sha256']==dec['candidate_manifest_sha256'] |
| if digest=='a52a68c8950c7eada39215d52b7b1eaf893d23577297951019bce382689baf48': |
| assert not allow_control |
| else: |
| certificate=enc_run['bundle_identity'] |
| assert certificate==dec['bundle_identity'] and certificate['candidate_manifest_sha256']==digest |
| assert certificate['trained_bundle'] and certificate['control_weights']==allow_control |
| assert certificate['weights_sha256']==manifest['files_sha256']['weights.safetensors'] |
| if manifest['format']=='shared-character-private-continuation-source-v1': |
| assert manifest['parent_bundle_sha256']=='b9281e7da68d2b9c331697a9fb15da3a2366fb937fce00742d658987d924318e' |
| assert manifest['root_bundle_sha256']=='a52a68c8950c7eada39215d52b7b1eaf893d23577297951019bce382689baf48' |
| assert manifest['parent_training_step']==8000 and manifest['cumulative_training_step']==8000+manifest['training_step'] |
| else: |
| assert manifest['format']=='shared-character-private-trained-source-v1' |
| assert manifest['parent_bundle_sha256']=='a52a68c8950c7eada39215d52b7b1eaf893d23577297951019bce382689baf48' |
| assert manifest['training_step'] in ([2] if allow_control else [2000,8000]) |
| assert sha(encoder/'report.json')==dec['encoder_report_sha256'] |
| assert enc['encoder_numerical_pass'] and enc['saved_neural_population_byte_exact'] |
| assert dec['total_graph_neural_values']==manifest['parameters']==74842752 |
| for name in ['character-assets/character_codec.json','english-assets/sentencepiece.bpe.model','english-assets/tokenizer_config.json']: |
| assert sha(candidate/name)==manifest['files_sha256'][name] |
| self.codec=CharacterCodec(list(read(candidate/'character-assets/character_codec.json')['character_ids'])) |
| assert self.codec.classes==573 |
| token_config=read(candidate/'english-assets/tokenizer_config.json') |
| assert token_config['tokenizer_class']=='XLMRobertaTokenizer' and token_config.get('clean_up_tokenization_spaces',False) is False |
| self.pieces=spm.SentencePieceProcessor(model_file=str(candidate/'english-assets/sentencepiece.bpe.model')) |
| assert self.pieces.get_piece_size()==64000 |
| library=dec['encoder_library'];resolved_library=str(library_path or library['path']) |
| assert sha(resolved_library)==library['sha256'] |
| self.operator_library_path=str(Path(resolved_library).resolve()) |
| self.graph_files=[] |
| global _LIBRARY_HASH |
| if _LIBRARY_HASH is None: |
| load_op_meta_info_and_register_op(resolved_library);_LIBRARY_HASH=library['sha256'] |
| else:assert _LIBRARY_HASH==library['sha256'],'Use a fresh process for a different operator library' |
|
|
| def load(folder,stem,hashes,selective): |
| for name,value in hashes.items():assert sha(folder/name)==value |
| self.graph_files.extend([str((folder/(stem+suffix)).resolve()) for suffix in ['.json','.pdiparams']]) |
| config=paddle.inference.Config(str(folder/(stem+'.json')),str(folder/(stem+'.pdiparams'))) |
| config.disable_gpu();config.set_cpu_math_library_num_threads(2) |
| if selective: |
| config.enable_mkldnn();config.enable_custom_passes(['ocr_conv_gelu_onednn_v1'],True) |
| else:config.disable_mkldnn() |
| config.switch_ir_optim(False);config.disable_glog_info() |
| return paddle.inference.create_predictor(config) |
| self.encoder=load(encoder,'encoder',enc['graph_sha256'],True) |
| self.decoders={branch:load(decoders/branch,'decoder-step',record['graph_sha256'],False) |
| for branch,record in dec['decoders'].items()} |
| assert set(self.encoder.get_input_names())=={'pixels','character_branch'} |
| assert all(set(p.get_input_names())=={'token','self_keys','self_values','cross_keys','cross_values'} for p in self.decoders.values()) |
|
|
| @staticmethod |
| def run(predictor,feeds): |
| assert set(predictor.get_input_names())==set(feeds) |
| for name,value in feeds.items():predictor.get_input_handle(name).copy_from_cpu(np.ascontiguousarray(value)) |
| predictor.run() |
| return [predictor.get_output_handle(n).copy_to_cpu() for n in predictor.get_output_names()] |
|
|
| def encode(self,pixels,branch): |
| assert branch in self.decoders and pixels.dtype==np.float32 and pixels.shape==(1,3,384,384) |
| outputs=self.run(self.encoder,dict(pixels=pixels,character_branch=np.array([branch=='character'],np.bool_))) |
| assert len(outputs)==2 and all(v.shape==(6,1,8,578,32) for v in outputs) |
| return [np.repeat(v,10,axis=1) for v in outputs] |
|
|
| def step(self,branch,token,self_keys,self_values,cross_keys,cross_values): |
| outputs=self.run(self.decoders[branch],dict(token=token,self_keys=self_keys,self_values=self_values, |
| cross_keys=cross_keys,cross_values=cross_values)) |
| assert len(outputs)==3 and outputs[0].shape==(10,64044 if branch=='english' else 573) |
| assert outputs[1].shape==outputs[2].shape==(6,10,8,self_keys.shape[3]+1,32) |
| assert all(v.dtype==np.float32 and np.isfinite(v).all() for v in outputs) |
| return outputs |
|
|
| def decode(self,tokens,branch): |
| if branch=='character':return self.codec.decode(tokens) |
| pieces=[] |
| for token in tokens: |
| if token in [0,1,2,3,64001]:continue |
| if not 4<=token<=64000:raise ValueError('Unassigned original output row') |
| pieces.append(self.pieces.id_to_piece(token-1)) |
| return ''.join(pieces).replace('\u2581',' ').strip() |
|
|
| def predict_pixels(self,pixels,branch): |
| cross=self.encode(pixels,branch) |
| state=TrOCRBeamState(beams=10,max_new_tokens=256 if branch=='english' else 257, |
| length_penalty=1.,suppress_tokens=() if branch=='english' else (0,1,3)) |
| keys=np.zeros((6,10,8,0,32),np.float32);values=keys.copy();token=np.full((10,1),2,np.int64) |
| while not state.done: |
| logits,next_keys,next_values=self.step(branch,token,keys,values,*cross) |
| parents,selected=state.advance(logits) |
| keys=np.ascontiguousarray(next_keys[:,parents]);values=np.ascontiguousarray(next_values[:,parents]);token=selected[:,None] |
| tokens,score=state.result() |
| return dict(tokens=tokens,text=self.decode(tokens,branch),score=score, |
| token_limit_reached=len(tokens)-1>=(256 if branch=='english' else 257) and tokens[-1]!=2, |
| decoder_iterations=state.current-1) |
|
|
