"""One model per process. Save each complete cell and its document predictions. Scores are strict character-boundary + type micro F1. No tuning on test labels. """ import argparse, contextlib, gc, hashlib, json, os, re, time, traceback from collections import defaultdict from pathlib import Path import numpy as np import torch from huggingface_hub import snapshot_download ROOT=Path(__file__).resolve().parents[1] def dump(path,obj): path.parent.mkdir(parents=True,exist_ok=True) tmp=path.with_suffix(path.suffix+'.tmp');tmp.write_text(json.dumps(obj,indent=2,ensure_ascii=False));tmp.replace(path) def chunks(text,size=64,overlap=16): words=list(re.finditer(r'\S+',text)) if not words:return [] out=[] for first in range(0,len(words),size-overlap): last=min(first+size,len(words));s=words[first].start();e=words[last-1].end() out.append((s,text[s:e])) if last==len(words):break return out def counts(gold,pred): g={tuple(x) for x in gold};p={tuple(x) for x in pred} return [len(g&p),len(p-g),len(g-p)] def metric(c): tp,fp,fn=map(int,c) return {'tp':tp,'fp':fp,'fn':fn,'precision':tp/(tp+fp) if tp+fp else 0., 'recall':tp/(tp+fn) if tp+fn else 0.,'f1':2*tp/(2*tp+fp+fn) if 2*tp+fp+fn else 0.} def summarize(rows): arr=np.array([r['counts'] for r in rows],dtype=np.int64).reshape(-1,3) out=metric(arr.sum(axis=0));out['documents']=len(rows) if len(rows)>1: rng=np.random.default_rng(20260928) boot=arr[rng.integers(0,len(rows),(1000,len(rows)))].sum(axis=1) den=2*boot[:,0]+boot[:,1]+boot[:,2] vals=np.divide(2*boot[:,0],den,out=np.zeros(len(boot)),where=den!=0) out['f1_95_ci_document_bootstrap']=[float(x) for x in np.quantile(vals,[.025,.975])] return out class Adapter: def __init__(self,entry): self.entry=entry;self.kind=entry['kind'];self.events=defaultdict(int) os.environ['TOKENIZERS_PARALLELISM']='false' torch.set_num_threads(4) path=snapshot_download(entry['id'],revision=entry['revision'],ignore_patterns=['*.bin','*.onnx','*.h5','*.msgpack','*.ot']) # Some T5 repositories contain only a PyTorch .bin checkpoint. if not list(Path(path).glob('*.safetensors')): path=snapshot_download(entry['id'],revision=entry['revision']) self.path=path if self.kind=='gliner2': from gliner2 import GLiNER2 self.model=GLiNER2.from_pretrained(path).to('cuda').eval();self.batch=8 self.max_len=512 if 'specialised' in entry['id'] else 3072 core=getattr(self.model,'model',self.model) proc=getattr(core,'processor',getattr(self.model,'processor',None)) self.runtime={'processor_max_length':getattr(proc,'max_length',None),'evaluation_max_len_argument':self.max_len} elif self.kind=='gliner': from gliner import GLiNER self.model=GLiNER.from_pretrained(path).to('cuda').eval();self.batch=8 self.runtime={'description_interface':'label: definition strings mapped back to canonical labels'} elif self.kind=='gner': from transformers import AutoTokenizer,AutoModelForSeq2SeqLM self.tokenizer=AutoTokenizer.from_pretrained(path) self.model=AutoModelForSeq2SeqLM.from_pretrained(path,dtype=torch.bfloat16).to('cuda').eval();self.batch=8 self.runtime={'precision':'bfloat16','generation':'greedy; max_new_tokens=512; official BIO prompt plus concise definitions'} elif self.kind=='nuextract': from transformers import AutoTokenizer,Qwen2_5_VLForConditionalGeneration self.tokenizer=AutoTokenizer.from_pretrained(path,padding_side='left') self.model=Qwen2_5_VLForConditionalGeneration.from_pretrained(path,dtype=torch.bfloat16).to('cuda').eval();self.batch=8 self.runtime={'precision':'bfloat16','generation':'greedy; max_new_tokens=512; native template with verbatim-string arrays'} else:raise ValueError(self.kind) self.runtime.update({'device':torch.cuda.get_device_name(),'parameter_count':sum(p.numel() for p in self.model.parameters())}) def predict(self,texts,schema): with torch.inference_mode(): if self.kind=='gliner2': raw=self.model.batch_extract_entities(texts,schema,batch_size=self.batch,threshold=.5,include_confidence=True,include_spans=True,max_len=self.max_len) return [[[x['start'],x['end'],label] for label,values in row.get('entities',{}).items() for x in values] for row in raw] if self.kind=='gliner': mapping={f'{k}: {v}':k for k,v in schema.items()} raw=self.model.batch_predict_entities(texts,list(mapping),threshold=.5,flat_ner=True,batch_size=self.batch) return [[[x['start'],x['end'],mapping.get(x['label'],x['label'])] for x in row] for row in raw] if self.kind=='gner': instruction=("Please analyze the sentence provided, identifying the type of entity for each word on a token-by-token basis.\n" "Output format is: word_1(label_1), word_2(label_2), ...\nWe'll use the BIO-format to label the entities, where:\n" "1. B- (Begin) indicates the start of a named entity.\n2. I- (Inside) is used for words within a named entity but are not the first word.\n" "3. O (Outside) denotes words that are not part of a named entity.\n") header=instruction+'\nUse the specific entity tags: '+', '.join(schema)+' and O.\nEntity definitions: '+json.dumps(schema)+'.\nSentence: ' encoded=self.tokenizer([header+t for t in texts],return_tensors='pt',padding=True,truncation=False).to('cuda') else: template=json.dumps({k:['verbatim-string'] for k in schema},ensure_ascii=False) prompts=[] for text in texts: # Definitions are metadata; only the delimited source is eligible for extraction. content='Entity definitions (instructions, not source text): '+json.dumps(schema)+'\nExtract only from the following source text:\n'+text messages=[{'role':'user','content':[{'type':'text','text':content}]}] prompts.append(self.tokenizer.apply_chat_template(messages,template=template,tokenize=False,add_generation_prompt=True)) encoded=self.tokenizer(prompts,return_tensors='pt',padding=True,truncation=False).to('cuda') outputs=self.model.generate(**encoded,max_new_tokens=512,do_sample=False) if self.kind=='nuextract':outputs=outputs[:,encoded['input_ids'].shape[1]:] result=[] for text,tokens in zip(texts,outputs): if len(tokens)>=512 and tokens[-1].item()!=self.tokenizer.eos_token_id:self.events['generation_limit_reached']+=1 response=self.tokenizer.decode(tokens,skip_special_tokens=True) if self.events['raw_output_samples_saved']<12: folder=ROOT/'diagnostics';folder.mkdir(exist_ok=True) with (folder/(self.entry['id'].replace('/','--')+'.jsonl')).open('a') as stream: stream.write(json.dumps({'text':text,'schema':schema,'output':response},ensure_ascii=False)+'\n') self.events['raw_output_samples_saved']+=1 if self.kind=='gner':result.append(self.parse_gner(text,response,schema)) else:result.append(self.parse_json(text,response,schema)) return result def parse_gner(self,text,response,schema): # Align BIO output to consecutive source tokens; retain unaligned entity predictions as FPs. # The author's model emits whitespace-separated tagged words even though # the prompt illustrates commas. Parse tags independently of separators. pairs=[];previous_end=0 for match in re.finditer(r'\((B-[^()]+|I-[^()]+|O)\)',response): word=response[previous_end:match.start()].strip() if word.startswith(', '):word=word[2:].strip() pairs.append((word,match.group(1)));previous_end=match.end() source=list(re.finditer(r'\S+',text));cursor=0;ents=[];active=None;unknown=0 aliases={k.casefold():k for k in schema} for word,tag in pairs: word=word.strip();prefix,_,label=tag.partition('-');label=aliases.get(label.casefold(),label) found=None for i in range(cursor,min(cursor+12,len(source))): if source[i].group()==word:found=i;break if found is not None: s,e=source[found].span();cursor=found+1 else:s,e=-1,-1;self.events['unaligned_output_tokens']+=1 if active and (prefix!='I' or label!=active[2] or s<0 or any(not c.isspace() for c in text[active[1]:s])): ents.append(active);active=None if prefix in ('B','I'): if s<0: unknown+=1;ents.append([-100000-unknown,-100000-unknown,label]);continue if active:active[1]=e else:active=[s,e,label] if active:ents.append(active) if not pairs:self.events['unparseable_outputs']+=1 return ents def parse_json(self,text,response,schema): response=response.strip();start=response.find('{');end=response.rfind('}') try:obj=json.loads(response[start:end+1]) except Exception: self.events['invalid_json_outputs']+=1 return [[-999999,-999999,'__invalid_output__']] result=[];used=set();bad=0 for label,values in obj.items(): if isinstance(values,str):values=[values] if not isinstance(values,list):self.events['invalid_value_types']+=1;continue for value in values: if not isinstance(value,str) or not value:continue found=next((m for m in re.finditer(re.escape(value),text) if (m.start(),m.end(),label) not in used),None) if found: used.add((found.start(),found.end(),label));result.append([found.start(),found.end(),label]) else: bad+=1;result.append([-100000-bad,-100000-bad,label]);self.events['unaligned_strings']+=1 return result def evaluate(adapter,name,meta,track,docs,exclude,deadline): root=ROOT/'results'/adapter.entry['id'].replace('/','--');path=root/f'{name}.{track}.json' if path.exists() and json.loads(path.read_text()).get('complete'):return decisions=meta['labels'] if track=='benchmark_schema': schema={k:v['definition'] for k,v in decisions.items()};mapping={k:k for k in schema} else: mapping={k:v['training_label'] for k,v in decisions.items() if v['status']=='familiar_equivalent'} schema={v['training_label']:v['training_definition'] for v in decisions.values() if v['status']=='familiar_equivalent'} if not schema: dump(path,{'complete':True,'not_applicable':True,'dataset':name,'track':track});return rows=[];start=time.time();model_seconds=0;total_chunks=0 for base in range(0,len(docs),8): if time.time()>deadline:raise TimeoutError('Evaluation deadline reached') batch=docs[base:base+8];parts=[];owner=[] for i,d in enumerate(batch): for offset,txt in chunks(d['text']):parts.append(txt);owner.append((i,offset)) predictions=[[] for _ in batch] for at in range(0,len(parts),8): tick=time.time() request=parts[at:at+8] try: outputs=adapter.predict(request,schema) if adapter.batch>1 else [adapter.predict([txt],schema)[0] for txt in request] except torch.cuda.OutOfMemoryError: torch.cuda.empty_cache();adapter.batch=1 outputs=[adapter.predict([txt],schema)[0] for txt in request] model_seconds+=time.time()-tick if len(outputs)!=len(request): outputs=[adapter.predict([txt],schema)[0] for txt in request] assert len(outputs)==len(request) for j,values in enumerate(outputs): i,offset=owner[at+j] for s,e,label in values: if s>=0: if not 0<=s