Spaces:
Sleeping
Sleeping
| """JPT-4B through its published llm2jev scoring interface on a local GPU.""" | |
| from functools import lru_cache | |
| import platform | |
| class AppleJPT: | |
| """Published llm2jev prompts and probabilities with MLX model execution.""" | |
| def __init__(self, model_name): | |
| from mlx_vlm import load | |
| from src.llm.model_paths import prepared_model_path | |
| print(f'[JPT] Loading {model_name} with MLX (original BF16 weights)',flush=True) | |
| self.model,self.processor=load(prepared_model_path(model_name)) | |
| self.prefix_tokens=[];self.prefix_cache=None | |
| def warm(self,prefix,images): | |
| if images:raise ValueError('The study uses text-only questions.') | |
| import mlx.core as mx | |
| # Leave a boundary margin so tokenization cannot merge the final prefix | |
| # token with the question. Every scoring request verifies this prefix. | |
| tokens=self.processor.tokenizer.encode(prefix,add_special_tokens=False)[:-2] | |
| if tokens==self.prefix_tokens:return | |
| self.prefix_tokens=tokens;lm=self.model.language_model | |
| self.prefix_cache=lm.make_cache() | |
| if tokens: | |
| lm(mx.array([tokens]),cache=self.prefix_cache,skip_logits=True) | |
| mx.eval([c.state for c in self.prefix_cache]) | |
| def score(self,text,images,ids): | |
| if images:raise ValueError('The study uses text-only questions.') | |
| rows,count=self._score_texts([text],ids) | |
| return rows[0],count | |
| def _score_texts(self,texts,ids): | |
| import mlx.core as mx | |
| tokens=[self.processor.tokenizer.encode(text,add_special_tokens=False) for text in texts] | |
| lm=self.model.language_model | |
| cache=lm.make_cache();suffixes=tokens | |
| if self.prefix_tokens and all(t[:len(self.prefix_tokens)]==self.prefix_tokens for t in tokens): | |
| from mlx.utils import tree_map | |
| for fresh,saved in zip(cache,self.prefix_cache): | |
| snapshot=tree_map(lambda x:mx.repeat(x,len(texts),axis=0) if isinstance(x,mx.array) else x, | |
| saved.prefix_cache_snapshot()) | |
| fresh.prefix_cache_restore(snapshot) | |
| suffixes=[t[len(self.prefix_tokens):] for t in tokens] | |
| lengths=[len(s) for s in suffixes];width=max(lengths) | |
| # Future padding cannot affect earlier positions in this causal model. | |
| # Select each row's final real token before projecting to vocabulary. | |
| padded=[s+[0]*(width-len(s)) for s in suffixes] | |
| out=lm(mx.array(padded),cache=cache,skip_logits=True,return_hidden=True) | |
| hidden=out.hidden_states[-1][mx.arange(len(texts)),mx.array(lengths)-1][:,None,:] | |
| logits=lm.speculative_logits_from_hidden(hidden)[:,0].astype(mx.float32) | |
| values=(logits-mx.logsumexp(logits,axis=-1,keepdims=True))[:,mx.array(ids)].tolist() | |
| return values,sum(map(len,tokens)) | |
| def score_many(self,texts,ids): | |
| rows=[];count=0 | |
| for i in range(0,len(texts),4): | |
| batch,n=self._score_texts(texts[i:i+4],ids);rows.extend(batch);count+=n | |
| print(f'[JPT] {min(i+4,len(texts))}/{len(texts)} questions scored',flush=True) | |
| return rows,count | |
| def load_jpt(model_name): | |
| from llm2jev.engine import LLM2Jev | |
| if platform.system()=='Darwin' and platform.machine()=='arm64': | |
| backend=AppleJPT(model_name) | |
| return LLM2Jev(backend.processor,backend,temperature=1.036,workers=1) | |
| import torch | |
| from llm2jev.backends import HF | |
| class LastTokenHF(HF): | |
| def __init__(self, model, device=None, **kwargs): | |
| # Loading safetensors directly onto MPS starts concurrent transfer | |
| # threads. Load on CPU and transfer serially to avoid a native | |
| # allocator/Metal crash in that loading path on this Mac. | |
| super().__init__(model, device='cpu' if device=='mps' else device, **kwargs) | |
| if device=='mps':self.model.to(device) | |
| # Restrict the LM projection to the last position. The upstream HF | |
| # reference projects every EDF token, wasting substantial GPU memory. | |
| def score(self, text, images, ids): | |
| if images: | |
| raise ValueError('This experiment uses text-only JPT inputs.') | |
| rows, count = self.score_many([text], ids) | |
| return rows[0], count | |
| def score_many(self, texts, ids): | |
| self.tok.padding_side = 'left' | |
| rows = []; count = 0 | |
| for start in range(0, len(texts), 4): | |
| inputs = self.tok(texts[start:start+4],padding=True,return_tensors='pt', | |
| add_special_tokens=False) | |
| count += int(inputs['attention_mask'].sum()) | |
| inputs = {k:v.to(self.model.device) for k,v in inputs.items()} | |
| with self.lock, torch.inference_mode(): | |
| logits = self.model(**inputs,logits_to_keep=1).logits[:,-1].float().log_softmax(-1) | |
| rows.extend(logits[:,ids].cpu().tolist()) | |
| return rows,count | |
| device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu' | |
| print(f'[JPT] Loading {model_name} on {device}',flush=True) | |
| from src.llm.model_paths import prepared_model_path | |
| backend = LastTokenHF(prepared_model_path(model_name),device=device,dtype='bfloat16') | |
| return LLM2Jev(backend.processor or backend.tok,backend,temperature=1.036,workers=1) | |
| def predict_jpt(state, questions, model_name): | |
| answers, usage = load_jpt(model_name).run(state, questions) | |
| return {'answers':answers,'model':model_name,'usage':usage, | |
| 'runtime':'llm2jev/mlx' if platform.system()=='Darwin' and platform.machine()=='arm64' else 'llm2jev/hf', | |
| 'temperature':1.036} | |