"""Full OCRBench v1, native Qwen vision classes and pinned upstream scoring semantics.

Upstream: MIT, Yuliang-Liu/MultimodalOCR @ 4e363b3f499811b43d792ef0bb33932ef05eca27.
The complete question/reference metadata must match the upstream 1,000-item list.
Inference budgets are our own pinned protocol; no vendor harness parity claim.
"""
from __future__ import annotations
import hashlib,importlib.metadata,io,json,platform,sys,time
from pathlib import Path

META_SHA='c1ccebb13035e4a37357fe3df4132152d674704d3df9bc07747f76a6cfd4c677'
RECIPE=dict(metadata_sha256=META_SHA,batch_size=4,max_new_tokens=128,min_pixels=200704,max_pixels=802816,do_sample=False,num_beams=1,dtype='bfloat16',attn_implementation='sdpa',prompt='Original question only, using the pinned model chat template',upstream_revision='4e363b3f499811b43d792ef0bb33932ef05eca27')

def score(prediction,answers,dataset):
    if not isinstance(prediction,str) or not isinstance(answers,list) or not all(isinstance(a,str) for a in answers):raise ValueError('Invalid OCRBench answer schema')
    def normalized(text):
        if dataset=='HME100k':return text.strip().replace('\n',' ').replace(' ','')
        return text.lower().strip().replace('\n',' ')
    response=normalized(prediction)
    return int(any(normalized(answer) in response for answer in answers))

def metadata_digest(rows):
    values=[{k:r[k] for k in ('dataset','question','question_type','answer')} for r in rows]
    return hashlib.sha256(json.dumps(values,ensure_ascii=False,sort_keys=True,separators=(',',':')).encode()).hexdigest()

def load_rows(task):
    import pyarrow.parquet as pq
    from huggingface_hub import hf_hub_download
    path=hf_hub_download(task['dataset'],'data/test-00000-of-00001.parquet',repo_type='dataset',revision=task['dataset_revision'])
    rows=pq.read_table(path,use_threads=False).to_pylist()
    if len(rows)!=1000 or metadata_digest(rows)!=META_SHA:raise ValueError('OCRBench mirror differs from the complete upstream item list')
    return rows

def evaluate(job,output):
    import torch
    from PIL import Image
    from transformers import AutoProcessor,AutoConfig,Qwen2VLForConditionalGeneration,Qwen2_5_VLForConditionalGeneration
    t=job['task'];model_id=job['model'];revision=job['revision']
    if t['recipe']!=RECIPE or t['limit']!=1000 or t['coverage']!='full' or t['metric']!='accuracy' or t['adapter']!='ocrbench':raise ValueError('Unsupported OCRBench protocol')
    if not any(m['id']==model_id and m['revision']==revision for m in t['models']):raise ValueError('Unpinned model')
    if not torch.cuda.is_available():raise RuntimeError('This protocol requires a CUDA GPU')
    torch.manual_seed(42);torch.set_num_threads(4);torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False
    started=time.monotonic();rows=load_rows(t);data_seconds=time.monotonic()-started;began=time.monotonic()
    config=AutoConfig.from_pretrained(model_id,revision=revision,trust_remote_code=False)
    cls={'qwen2_vl':Qwen2VLForConditionalGeneration,'qwen2_5_vl':Qwen2_5_VLForConditionalGeneration}.get(config.model_type)
    if cls is None:raise ValueError('Model requires a different reviewed native adapter')
    processor=AutoProcessor.from_pretrained(model_id,revision=revision,trust_remote_code=False,min_pixels=RECIPE['min_pixels'],max_pixels=RECIPE['max_pixels'],use_fast=False)
    processor.tokenizer.padding_side='left'
    model=cls.from_pretrained(model_id,revision=revision,trust_remote_code=False,torch_dtype=torch.bfloat16,attn_implementation='sdpa',low_cpu_mem_usage=True).to('cuda:0').eval()
    parameters=sum(p.numel() for p in model.parameters());load_seconds=time.monotonic()-began;samples=[];began=time.monotonic()
    # A diagnostic mode never emits a publishable result; it does not alter the full protocol.
    diagnostic=bool(job.get('diagnostic_only'));diagnostic_indices=(0,49,300,500,700,899,900,999)
    selected=[rows[i] for i in diagnostic_indices] if diagnostic else rows
    for offset in range(0,len(selected),4):
        group=selected[offset:offset+4];images=[];texts=[];sizes=[];hashes=[]
        for row in group:
            blob=row['image'];blob=blob['bytes'] if isinstance(blob,dict) else None
            if not isinstance(blob,bytes):raise ValueError('Missing pinned image bytes')
            with Image.open(io.BytesIO(blob)) as image:rgb=image.convert('RGB')
            images.append(rgb);sizes.append(list(rgb.size));hashes.append(hashlib.sha256(rgb.tobytes()).hexdigest())
            messages=[dict(role='user',content=[dict(type='image'),dict(type='text',text=row['question'])])]
            texts.append(processor.apply_chat_template(messages,tokenize=False,add_generation_prompt=True))
        inputs=processor(text=texts,images=images,padding=True,return_tensors='pt').to('cuda:0')
        for key,value in inputs.items():
            if torch.is_floating_point(value):inputs[key]=value.to(torch.bfloat16)
        tick=time.monotonic()
        with torch.inference_mode():tokens=model.generate(**inputs,max_new_tokens=128,do_sample=False,num_beams=1)
        generated=tokens[:,inputs['input_ids'].shape[1]:];predictions=processor.batch_decode(generated,skip_special_tokens=True,clean_up_tokenization_spaces=False)
        torch.cuda.synchronize();elapsed=time.monotonic()-tick
        for index,(row,prediction) in enumerate(zip(group,predictions,strict=True)):
            n=offset+index
            samples.append(dict(id=str(diagnostic_indices[n] if diagnostic else n),question=row['question'],reference=row['answer'],prediction=prediction,score=score(prediction,row['answer'],row['dataset']),dataset=row['dataset'],question_type=row['question_type'],image_size=sizes[index],image_rgb_sha256=hashes[index],batch_seconds=elapsed,batch_size_actual=len(group)))
        print('Scored',len(samples),'/',len(selected),flush=True)
    seconds=time.monotonic()-began;groups={}
    for sample in samples:
        g=groups.setdefault(sample['question_type'],dict(correct=0,n=0));g['correct']+=sample['score'];g['n']+=1
    if diagnostic:
        output.write_text(json.dumps(dict(diagnostic_only=True,n=len(samples),samples=samples,seconds=seconds,peak_gpu_bytes=torch.cuda.max_memory_allocated())));return
    result=dict(score=sum(s['score'] for s in samples)/len(samples),n=len(samples),model_revision=revision,dataset_revision=t['dataset_revision'],runner_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),seconds=time.monotonic()-started,data_loading_seconds=data_seconds,model_load_seconds=load_seconds,scoring_seconds=seconds,peak_gpu_bytes=torch.cuda.max_memory_allocated(),parameters=parameters,protocol=t,seed=42,dtype='bfloat16',tf32=False,batch_size=4,subgroups=groups,upstream_total_score=sum(s['score'] for s in samples),upstream_max_score=1000,metadata_sha256=META_SHA,notes='All 1,000 OCRBench v1 questions. Upstream answer-containment scoring; response verbosity can game this metric. Separate visual-text QA protocol, not document parsing/layout or table reconstruction. Image/decoding budgets and model chat template are CodeSOTA-pinned, not a claim of vendor inference parity.')
    env=dict(python=sys.version,platform=platform.platform(),torch=torch.__version__,cuda=torch.version.cuda,gpu=torch.cuda.get_device_name(),packages={d.metadata['Name']:d.version for d in importlib.metadata.distributions()})
    output.parent.mkdir(parents=True,exist_ok=True);output.write_text(json.dumps(dict(result=result,samples=samples,environment=env),allow_nan=False))

if __name__=='__main__':evaluate(json.loads(Path(sys.argv[1]).read_text()),Path(sys.argv[2]))