Spaces:
Running on Zero
Running on Zero
| """Sparse, smoothed trigram corrections and exact second-order chain decoding.""" | |
| from pathlib import Path | |
| import math | |
| import torch | |
| MAX_TRIGRAM_LOOKUP_ELEMENTS = 2_000_000 | |
| class TrigramScorer(torch.nn.Module): | |
| """Correction log P(c|a,b) - log P(c|b), with bigram backoff. | |
| All context totals precede singleton filtering during training. Unseen or | |
| low-count triples give zero correction; clipping limits corpus preferences. | |
| """ | |
| def __init__(self,vocab_size,triple_keys,triple_counts,prefix_keys,prefix_counts,*,alpha=50.,metadata=None): | |
| super().__init__() | |
| if vocab_size<1 or vocab_size**3>=2**63 or not math.isfinite(alpha) or alpha<=0: | |
| raise ValueError('Invalid trigram vocabulary or smoothing') | |
| self.vocab_size=int(vocab_size);self.alpha=float(alpha);self.metadata=metadata or {} | |
| for name,keys,counts,limit in [('triple',triple_keys,triple_counts,vocab_size**3),('prefix',prefix_keys,prefix_counts,vocab_size**2)]: | |
| keys=torch.as_tensor(keys,dtype=torch.long);counts=torch.as_tensor(counts,dtype=torch.float32) | |
| if keys.ndim!=1 or counts.shape!=keys.shape or (keys<0).any() or (keys>=limit).any() or (keys[1:]<=keys[:-1]).any(): | |
| raise ValueError('Invalid sorted sparse trigram keys') | |
| if not torch.isfinite(counts).all() or (counts<=0).any():raise ValueError('Invalid trigram counts') | |
| self.register_buffer(name+'_keys',keys);self.register_buffer(name+'_counts',counts) | |
| def _lookup(keys,counts,query): | |
| if not keys.numel():return torch.zeros_like(query,dtype=torch.float32) | |
| index=torch.searchsorted(keys,query.contiguous());safe=index.clamp_max(len(keys)-1) | |
| return torch.where((index<len(keys))&(keys[safe]==query),counts[safe],0.) | |
| def correction(self,a,b,c,bigram,*,min_count=2,clip=2.): | |
| if bigram.vocab_size!=self.vocab_size:raise ValueError('Bigram/trigram vocabulary mismatch') | |
| if not isinstance(min_count,int) or min_count<1 or not math.isfinite(clip) or clip<0: | |
| raise ValueError('Invalid trigram minimum count or clipping') | |
| if min_count < self.metadata.get('storage_min_count', 1): | |
| raise ValueError('Requested count threshold is below checkpoint storage threshold') | |
| a,b,c=torch.broadcast_tensors(a,b,c) | |
| for ids in (a,b,c): | |
| if ids.device!=self.triple_keys.device:raise ValueError('Move trigram scorer to candidate device') | |
| if ((ids<0)|(ids>=self.vocab_size)).any():raise ValueError('Token outside trigram vocabulary') | |
| context=a*self.vocab_size+b | |
| counts=self._lookup(self.triple_keys,self.triple_counts,context*self.vocab_size+c) | |
| totals=self._lookup(self.prefix_keys,self.prefix_counts,context) | |
| base=bigram.log_prob(b,c) | |
| smoothed=torch.logaddexp(counts.log(),math.log(self.alpha)+base)-(totals+self.alpha).log() | |
| return torch.where(counts>=min_count,(smoothed-base).clamp(-clip,clip),0.) | |
| def lattice_corrections(self,ids,bigram,*,min_count=2,clip=2.): | |
| k=ids.shape[1] | |
| if len(ids)<3:return torch.zeros((0,k,k,k),device=ids.device) | |
| if bigram.vocab_size!=self.vocab_size:raise ValueError('Bigram/trigram vocabulary mismatch') | |
| if not isinstance(min_count,int) or min_count<1 or not math.isfinite(clip) or clip<0: | |
| raise ValueError('Invalid trigram minimum count or clipping') | |
| if min_count<self.metadata.get('storage_min_count',1): | |
| raise ValueError('Requested count threshold is below checkpoint storage threshold') | |
| if ids.device!=self.triple_keys.device:raise ValueError('Move trigram scorer to candidate device') | |
| if ((ids<0)|(ids>=self.vocab_size)).any():raise ValueError('Token outside trigram vocabulary') | |
| # At ordinary K=16 this is one vectorized lookup across all positions. | |
| # Large K uses window chunks so sparse-lookup temporaries do not grow | |
| # into multiple gigabytes beyond the required K^3 potential table. | |
| chunk=max(1,MAX_TRIGRAM_LOOKUP_ELEMENTS//(k**3)) | |
| output=torch.empty((len(ids)-2,k,k,k),device=ids.device,dtype=torch.float32) | |
| for start in range(0,len(ids)-2,chunk): | |
| end=min(len(ids)-2,start+chunk) | |
| # Prefix totals depend on (a,b), and backoff on (b,c). Looking them | |
| # up on the full K^3 broadcast repeats each sparse search K times. | |
| a,b,c=ids[start:end],ids[start+1:end+1],ids[start+2:end+2] | |
| context=a[:,:,None]*self.vocab_size+b[:,None,:] | |
| totals=self._lookup(self.prefix_keys,self.prefix_counts,context)[:,:,:,None] | |
| base=bigram.log_prob(b[:,:,None],c[:,None,:])[:,None,:,:] | |
| counts=self._lookup(self.triple_keys,self.triple_counts, | |
| context[:,:,:,None]*self.vocab_size+c[:,None,None,:]) | |
| smoothed=torch.logaddexp(counts.log(),math.log(self.alpha)+base)-(totals+self.alpha).log() | |
| output[start:end]=torch.where(counts>=min_count,(smoothed-base).clamp(-clip,clip),0.) | |
| return output | |
| def load(cls,path,*,bigram,device='cpu'): | |
| state=torch.load(Path(path),map_location='cpu',weights_only=True) | |
| if state.pop('version',None)!=1:raise ValueError('Unsupported trigram checkpoint') | |
| scorer=cls(**state) | |
| if scorer.vocab_size!=bigram.vocab_size or scorer.metadata.get('tokenizer_sha256')!=bigram.metadata.get('tokenizer_sha256'): | |
| raise ValueError('Bigram and trigram tokenizer mismatch') | |
| if scorer.metadata.get('tokens')!=bigram.metadata.get('tokens'): | |
| raise ValueError('Bigram and trigram training corpora mismatch') | |
| return scorer.to(device) | |
| def _potentials(emissions,edges,triples,temperature=1.): | |
| length,k=emissions.shape | |
| if edges.shape!=(length-1,k,k) or triples.shape!=(max(0,length-2),k,k,k): | |
| raise ValueError('Invalid second-order lattice shapes') | |
| if length<1 or k<1:raise ValueError('Empty second-order lattice') | |
| temperature=float(temperature) | |
| if not math.isfinite(temperature) or temperature<=0:raise ValueError('Temperature must be positive') | |
| scale=max(temperature,1e-6) | |
| unary=emissions.float();unary=(unary-unary.amax(-1,keepdim=True))/scale | |
| edges=edges.float();triples=triples.float() | |
| if edges.numel():edges=(edges-edges.amax((1,2),keepdim=True))/scale | |
| if triples.numel():triples=(triples-triples.amax((1,2,3),keepdim=True))/scale | |
| return unary,edges,triples | |
| def _forward(unary,edges,triples,*,maximum=False): | |
| length,k=unary.shape | |
| rows=[];back=[] | |
| if length<2:return rows,back | |
| row=unary[0,:,None]+unary[1,None,:]+edges[0] | |
| rows.append(row-row.max()) | |
| for t in range(2,length): | |
| score=rows[-1][:,:,None]+triples[t-2] | |
| if maximum: | |
| row,index=score.max(0);back.append(index) | |
| else:row=torch.logsumexp(score,0) | |
| row=row+edges[t-1]+unary[t,None,:] | |
| rows.append(row-row.max()) | |
| return rows,back | |
| def decode_second_order(ids,emissions,edges,triples): | |
| # A small chain launches hundreds of tiny GPU reductions. A single host | |
| # transfer and vectorized CPU reductions can cost less than those launches. | |
| # Keep the GPU route for large lattices, and sampling on its existing route. | |
| if ids.device.type=='cuda' and len(ids)<=256 and ids.shape[1]<=16: | |
| return _decode_second_order_host(ids,emissions,edges,triples) | |
| unary,edges,triples=_potentials(emissions,edges,triples) | |
| if len(ids)==1:return ids.gather(1,unary.argmax(-1,keepdim=True)).squeeze(1) | |
| forward,back=_forward(unary,edges,triples,maximum=True) | |
| k=ids.shape[1] | |
| if ids.device.type!='cuda': | |
| # Scalar indexing is inexpensive on CPU; retain its lean traceback. | |
| path=torch.empty(len(ids),device=ids.device,dtype=torch.long) | |
| last=forward[-1].reshape(-1).argmax();path[-2]=last//k;path[-1]=last%k | |
| for t in range(len(ids)-3,-1,-1):path[t]=back[t][path[t+1],path[t+2]] | |
| else: | |
| path=_traceback_device(back,forward[-1].reshape(-1).argmax(),k) | |
| return ids.gather(1,path[:,None]).squeeze(1) | |
| def _decode_second_order_host(ids,emissions,edges,triples): | |
| """Exact MAP recurrence on CPU, preserving float32 operations/tie order.""" | |
| import numpy as np | |
| # Normalize on the original device to retain the original score arithmetic. | |
| unary,edges,triples=_potentials(emissions,edges,triples) | |
| length,k=unary.shape | |
| # One synchronous copy instead of one per predecessor or lattice row. | |
| packed=torch.cat((unary.reshape(-1),edges.reshape(-1),triples.reshape(-1))) | |
| host=packed.detach().cpu().numpy() | |
| cut=length*k;end=cut+(length-1)*k*k | |
| u=host[:cut].reshape(length,k) | |
| e=host[cut:end].reshape(length-1,k,k) | |
| tri=host[end:].reshape(max(0,length-2),k,k,k) | |
| if length==1: | |
| path=np.array([u[0].argmax()],dtype=np.int64) | |
| else: | |
| row=u[0,:,None]+u[1,None,:]+e[0] | |
| row=row-row.max();back=[] | |
| for t in range(2,length): | |
| score=row[:,:,None]+tri[t-2] | |
| index=score.argmax(axis=0) | |
| row=score.max(axis=0)+e[t-1]+u[t,None,:] | |
| row=row-row.max();back.append(index) | |
| path=np.empty(length,dtype=np.int64) | |
| last=row.reshape(-1).argmax();path[-2]=last//k;path[-1]=last%k | |
| for t in range(length-3,-1,-1): | |
| path[t]=back[t][path[t+1],path[t+2]] | |
| indices=torch.from_numpy(path).to(device=ids.device) | |
| return ids.gather(1,indices[:,None]).squeeze(1) | |
| def _traceback_device(back,last,k): | |
| """Trace pair states without converting GPU indices to host scalars.""" | |
| pair=last.reshape(1) | |
| reverse=[(pair%k).squeeze(0),(pair//k).squeeze(0)] | |
| for predecessors in reversed(back): | |
| # A one-element tensor gather keeps the index on the GPU. Scalar | |
| # tensor indexing can otherwise synchronize with the host each hop. | |
| previous=predecessors.reshape(-1).gather(0,pair) | |
| reverse.append(previous.squeeze(0)) | |
| pair=previous*k+pair//k | |
| return torch.stack(reverse[::-1]) | |
| def sample_second_order(ids,emissions,edges,triples,temperature=1.): | |
| if temperature==0:return decode_second_order(ids,emissions,edges,triples) | |
| unary,edges,triples=_potentials(emissions,edges,triples,temperature) | |
| if len(ids)==1: | |
| index=torch.multinomial(unary[0].softmax(-1),1) | |
| return ids[0,index] | |
| forward,_=_forward(unary,edges,triples) | |
| k=ids.shape[1];path=torch.empty(len(ids),device=ids.device,dtype=torch.long) | |
| last=torch.multinomial(forward[-1].reshape(-1).softmax(-1),1).squeeze(0) | |
| path[-2]=last//k;path[-1]=last%k | |
| for t in range(len(ids)-3,-1,-1): | |
| scores=forward[t][:,path[t+1]]+triples[t,:,path[t+1],path[t+2]] | |
| path[t]=torch.multinomial(scores.softmax(-1),1).squeeze(0) | |
| return ids.gather(1,path[:,None]).squeeze(1) | |
| def second_order_marginals(emissions,edges,triples,temperature=1.,*,maximum=False): | |
| """Return pair log-potentials (max) or normalized joint pair probabilities.""" | |
| unary,edges,triples=_potentials(emissions,edges,triples,temperature or 1.) | |
| if len(unary)<2:return edges.clone() | |
| forward,_=_forward(unary,edges,triples,maximum=maximum) | |
| backward=[None]*(len(unary)-1);backward[-1]=torch.zeros_like(edges[-1]) | |
| for t in range(len(unary)-3,-1,-1): | |
| scores=triples[t]+edges[t+1][None,:,:]+unary[t+2,None,None,:]+backward[t+1][None,:,:] | |
| row=scores.max(-1).values if maximum else torch.logsumexp(scores,-1) | |
| backward[t]=row-row.max() | |
| pair_scores=torch.stack([a+b for a,b in zip(forward,backward)]) | |
| return pair_scores if maximum else pair_scores.flatten(1).softmax(-1).reshape_as(pair_scores) | |
| def second_order_path_margins(ids,emissions,edges,triples,selected): | |
| if len(ids)==1: | |
| from .bigram import BigramViterbiSelector | |
| return BigramViterbiSelector.path_margins(ids,emissions,edges,selected,1.) | |
| pairs=second_order_marginals(emissions,edges,triples,maximum=True) | |
| scores=torch.cat((pairs[0].max(1).values[None,:],pairs.max(1).values),0) | |
| same=ids==selected[:,None] | |
| return (scores.masked_fill(~same,-torch.inf).max(-1).values-scores.masked_fill(same,-torch.inf).max(-1).values).clamp_min(0.) | |