Spaces:
Sleeping
Sleeping
File size: 5,763 Bytes
c2b7494 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """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
@lru_cache(maxsize=1)
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}
|