File size: 8,519 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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """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)
|