| """Frozen FP32 ESM2 base encoder, fixed separate terminal windows.""" |
| import os |
| os.environ['HF_HUB_OFFLINE']='1';os.environ['TRANSFORMERS_OFFLINE']='1';os.environ['TOKENIZERS_PARALLELISM']='false' |
| import time,json,hashlib |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from transformers import EsmModel,EsmTokenizer |
| POLICY={'whole_sequence_max':512,'long_head':256,'long_tail':256,'long_windows_separate':True,'middle_omitted':True,'window_pooling':'mean final-layer non-special residue states','sequence_pooling':'one whole window or arithmetic mean of two terminal-window pools','dtype':'float32','CPU_batch_windows':4,'CPU_intraop_threads':2,'CPU_interop_threads':2,'unknown':'per-character unk mapped, special tokens excluded, zero if none','deduplicate_exact_sequences':True,'no_finetuning':True} |
| class FrozenEncoder: |
| def __init__(self,assets,device='cpu'): |
| torch.set_num_threads(2);torch.manual_seed(42);torch.use_deterministic_algorithms(True) |
| self.device=torch.device(device);assets=Path(assets) |
| self.tokenizer=EsmTokenizer.from_pretrained(assets,local_files_only=True) |
| self.model,info=EsmModel.from_pretrained(assets,local_files_only=True,use_safetensors=True,add_pooling_layer=False,output_loading_info=True) |
| assert not info['missing_keys'] and not info.get('mismatched_keys'),info |
| self.model.eval().requires_grad_(False).to(self.device);self.dimension=self.model.config.hidden_size;self.vocab=self.tokenizer.get_vocab();self.special=set(self.tokenizer.all_special_ids) |
| self.audit={'XBU_tokens':self.tokenizer.tokenize('AXBU'),'XBU_ids':{c:self.tokenizer.convert_tokens_to_ids(c) for c in 'XBU'},'special_ids':sorted(self.special),'loading_info':info,'last_hidden_state_not_logits':True,'model_class':type(self.model).__name__,'parameter_count':sum(p.numel() for p in self.model.parameters()),'eval_mode':not self.model.training,'all_parameters_frozen':all(not p.requires_grad for p in self.model.parameters())} |
| assert self.audit['XBU_tokens']==list('AXBU') and all(x not in self.special and x!=self.tokenizer.unk_token_id for x in self.audit['XBU_ids'].values()) |
| self.timing={} |
| def encode(self,sequences): |
| sequences=list(sequences) |
| if not sequences:return np.empty((0,self.dimension),np.float32) |
| assert all(isinstance(s,str) and s for s in sequences);unique=sorted(set(sequences));windows=[] |
| for i,s in enumerate(unique): |
| if len(s)<=512:windows.append((i,0,s)) |
| else:windows.extend([(i,0,s[:256]),(i,1,s[-256:])]) |
| windows.sort(key=lambda w:(len(w[2]),w[0],w[1]));acc=np.zeros((len(unique),self.dimension),np.float64);counts=np.zeros(len(unique),np.int32);toksecs=0.;forwardsecs=0.;poolsecs=0.;start=time.perf_counter() |
| with torch.inference_mode(): |
| for offset in range(0,len(windows),4): |
| batch=windows[offset:offset+4];t=time.perf_counter();width=max(len(x[2]) for x in batch)+2 |
| ids=torch.full((len(batch),width),self.tokenizer.pad_token_id,dtype=torch.long);attention=torch.zeros_like(ids);residues=torch.zeros_like(ids,dtype=torch.bool) |
| for j,(_,_,s) in enumerate(batch): |
| tokens=[self.vocab.get(c,self.tokenizer.unk_token_id) for c in s] |
| ids[j,:len(s)+2]=torch.tensor([self.tokenizer.cls_token_id]+tokens+[self.tokenizer.eos_token_id]);attention[j,:len(s)+2]=1;residues[j,1:len(s)+1]=torch.tensor([c not in self.special for c in tokens]) |
| ids,attention,residues=ids.to(self.device),attention.to(self.device),residues.to(self.device);toksecs+=time.perf_counter()-t;t=time.perf_counter() |
| hidden=self.model(input_ids=ids,attention_mask=attention,return_dict=True).last_hidden_state |
| if self.device.type=='cuda':torch.cuda.synchronize() |
| forwardsecs+=time.perf_counter()-t;t=time.perf_counter();denom=residues.sum(dim=1).clamp_min(1).unsqueeze(1);pooled=(hidden*residues.unsqueeze(-1)).sum(dim=1)/denom;values=pooled.cpu().numpy().astype(np.float32) |
| for j,(i,_,_) in enumerate(batch):acc[i]+=values[j].astype(np.float64);counts[i]+=1 |
| poolsecs+=time.perf_counter()-t |
| assert counts.tolist()==[1 if len(s)<=512 else 2 for s in unique];values=(acc/counts[:,None]).astype(np.float32);assert np.isfinite(values).all();mapping=dict(zip(unique,values));out=np.stack([mapping[s] for s in sequences]) |
| self.timing={'unique_sequences':len(unique),'windows':len(windows),'window_residues':sum(len(x[2]) for x in windows),'tokenization_seconds':toksecs,'forward_seconds':forwardsecs,'pooling_alignment_seconds':poolsecs,'encoding_seconds':time.perf_counter()-start};return out |
| def extraction_hash(revision):return hashlib.sha256(json.dumps({'revision':revision,'policy':POLICY},sort_keys=True).encode()).hexdigest() |
|
|
