"""Live BPB v1: pinned raw UTF-8 documents; all content tokens scored once, no chat template."""
import hashlib,importlib.metadata,json,math,os,platform,sys,time
from pathlib import Path

def windows(tokens,start_id,context=1024,stride=512):
 ids=[start_id]+tokens;cursor=1
 while cursor<len(ids):
  end=min(cursor+stride,len(ids));begin=max(0,end-context)
  yield ids[begin:end],ids[cursor:end]
  cursor=end

def evaluate(model_info,documents,output):
 import torch
 from transformers import AutoTokenizer,AutoModelForCausalLM
 torch.manual_seed(42);torch.set_num_threads(4);torch.backends.cuda.matmul.allow_tf32=False
 began=time.monotonic();tokenizer=AutoTokenizer.from_pretrained(model_info['id'],revision=model_info['revision'],trust_remote_code=False)
 model=AutoModelForCausalLM.from_pretrained(model_info['id'],revision=model_info['revision'],trust_remote_code=False,torch_dtype=torch.float32).to('cuda:0').eval();torch.cuda.synchronize();load_seconds=time.monotonic()-began
 parameters=sum(p.numel() for p in model.parameters())
 if not 0<parameters<1_000_000_000:raise ValueError('Model outside strict under-1B cohort')
 start_id=tokenizer.bos_token_id if tokenizer.bos_token_id is not None else tokenizer.eos_token_id
 if start_id is None:raise ValueError('No explicit start-token policy available')
 examples=[];per_doc=[]
 for doc in documents:
  text=doc['text'];raw=text.encode('utf-8')
  if hashlib.sha256(raw).hexdigest()!=doc['sha256'] or len(raw)!=doc['bytes']:raise ValueError('Corpus integrity mismatch')
  tokens=tokenizer.encode(text,add_special_tokens=False)
  if not tokens or not raw:raise ValueError('Empty document')
  per_doc.append(dict(id=doc['id'],published=doc['published'],sha256=doc['sha256'],bytes=len(raw),tokens=len(tokens),nll_nats=0.0))
  for ids,target in windows(tokens,start_id):examples.append((len(per_doc)-1,ids,target))
 torch.cuda.reset_peak_memory_stats();torch.cuda.synchronize();began=time.monotonic()
 for offset in range(0,len(examples),4):
  batch=examples[offset:offset+4];length=max(len(ids) for _,ids,_ in batch);keep=max(len(target) for _,_,target in batch)+1
  ids=torch.full((len(batch),length),start_id,device='cuda:0',dtype=torch.long);mask=torch.zeros_like(ids)
  for row,(_,values,_) in enumerate(batch):ids[row,-len(values):]=torch.tensor(values,device='cuda:0');mask[row,-len(values):]=1
  with torch.inference_mode():
   logits=model(ids,attention_mask=mask,position_ids=(mask.cumsum(-1)-1).clamp_min(0),use_cache=False,logits_to_keep=keep).logits
   values=[]
   for row,(_,_,target) in enumerate(batch):
    selected=logits[row,-len(target)-1:-1].float();target_ids=torch.tensor(target,device='cuda:0')
    values.append(-torch.log_softmax(selected,dim=-1).gather(1,target_ids[:,None]).double().sum())
   losses=torch.stack(values).cpu().tolist()
  for (doc,_,_),loss in zip(batch,losses):per_doc[doc]['nll_nats']+=loss
 torch.cuda.synchronize();scoring_seconds=time.monotonic()-began
 byte_count=sum(d['bytes'] for d in per_doc);token_count=sum(d['tokens'] for d in per_doc)
 manifest=json.loads(Path('/var/lib/codesota-sub1b/corpus-manifest.json').read_text())
 protocol=dict(id='codesota-live-bpb-v1',task='language-modeling',benchmark=manifest['id'],metric='bits_per_utf8_byte',direction='min',context_tokens=1024,stride_tokens=512,batch_size=4,dtype='float32',tf32=False,start_token_id=start_id,start_token_policy='tokenizer BOS, else EOS; prefix not scored',tokenization='raw text, no chat template or added special tokens; every content token scored once',aggregation='sum document negative log likelihood / log(2) / sum raw UTF-8 bytes',selection='20 newest cs.CL abstracts in pinned temporal manifest; domain-specific pilot, not representative general text',runner_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest())
 result=dict(id=hashlib.sha256((model_info['revision']+protocol['runner_sha256']+json.dumps(per_doc,sort_keys=True)).encode()).hexdigest()[:24],system=model_info,parameters=parameters,track='open',origin='codesota',scope='subset',protocol=protocol,score=sum(d['nll_nats'] for d in per_doc)/math.log(2)/byte_count,documents=len(per_doc),bytes=byte_count,tokens=token_count,model_load_seconds=load_seconds,scoring_seconds=scoring_seconds,scoring_tokens_per_second=token_count/scoring_seconds,gpu_slot_seconds=load_seconds+scoring_seconds,peak_gpu_bytes=torch.cuda.max_memory_allocated(),finished=int(time.time()),per_document=per_doc,environment=dict(python=sys.version,torch=torch.__version__,transformers=__import__('transformers').__version__,gpu=torch.cuda.get_device_name(),platform=platform.platform()))
 output.parent.mkdir(parents=True,exist_ok=True);output.write_text(json.dumps(result,indent=2,allow_nan=False));print(model_info['id'],'BPB',round(result['score'],4),'parameters',parameters,flush=True)
if __name__=='__main__':
 root=Path(__file__).resolve().parent;models=json.loads((root/'sub1b-models.json').read_text());docs=json.loads(Path('/var/lib/codesota-sub1b/corpus.json').read_text())
 for model in models:evaluate(model,docs,root/'state'/'sub1b'/model['id'].split('/')[-1]/'result.json')