| """Bounded, inspectable span rewriting with exact preservation of surrounding text.""" |
| import math |
| import time |
| import torch |
| from torch.nn import functional as F |
| from .model import BOS,EOS,MASK |
| from .inference import denoise |
| from .train import amp |
|
|
|
|
| @torch.inference_mode() |
| def local_nll(model,tokenizer,text,start,end,precision="fp32"): |
| encoded=tokenizer.encode(text) |
| if len(encoded.ids)+1>model.config.max_length: |
| raise ValueError("Text exceeds model context") |
| positions=[i for i,(a,b) in enumerate(encoded.offsets) if b>start and a<end] |
| if not positions:raise ValueError("Selection has no tokens") |
| last=min(len(encoded.ids),positions[-1]+1+16) |
| device=next(model.parameters()).device |
| tokens=torch.tensor([[BOS]+encoded.ids],device=device) |
| targets=torch.tensor(encoded.ids+[EOS],device=device) |
| selected=torch.zeros_like(tokens,dtype=torch.bool) |
| selected[0,positions[0]:last]=True |
| with amp(device.type,precision): |
| logits=model(tokens,torch.zeros(1,device=device),selected,causal=True,right_padded=True) |
| return float(F.cross_entropy(logits.float(),targets[selected[0]])) |
|
|
|
|
| @torch.inference_mode() |
| def rewrite_span(model,tokenizer,text,start,end,seed=42,precision="fp32",ranker=None): |
| if not isinstance(text,str) or not 1<=len(text)<=3000: |
| raise ValueError("Use a short English text (1..3000 characters)") |
| if not isinstance(start,int) or not isinstance(end,int) or not 0<=start<end<=len(text): |
| raise ValueError("Select a passage inside the text") |
| if not isinstance(seed,int) or not 0<=seed<2**63-16: |
| raise ValueError("Invalid seed") |
| original=text[start:end] |
| a=start+len(original)-len(original.lstrip()) |
| b=end-(len(original)-len(original.rstrip())) |
| if a>=b or len(text[a:b].split())>12: |
| raise ValueError("Select 1 to 12 words") |
| if (a>0 and text[a-1].isalnum() and text[a].isalnum()) or (b<len(text) and text[b-1].isalnum() and text[b].isalnum()): |
| raise ValueError("Select complete words") |
| if len(tokenizer.encode(text).ids)>model.config.max_length-32: |
| raise ValueError("Text too long; shorten it before editing") |
| prefix,suffix=text[:a],text[b:] |
| left=tokenizer.encode(prefix.rstrip()).ids |
| right=tokenizer.encode(suffix).ids |
| count=len(tokenizer.encode((' ' if prefix else '')+text[a:b]).ids) |
| if not 1<=count<=16:raise ValueError("Select a shorter passage (at most 16 subword tokens)") |
| before=local_nll(model,tokenizer,text,a,b,precision) |
| candidates=[];seen={text[a:b]};started=time.perf_counter() |
| |
| lengths=[count,count,count,max(1,count-1),min(16,count+1),count] |
| for attempt,length in enumerate(lengths): |
| ids=[BOS]+left+[MASK]*length+right+[EOS] |
| result=denoise(model,tokenizer,ids,steps=12,temperature=0 if attempt==0 else .8, |
| seed=seed+attempt,precision=precision,allow_eos=False) |
| fragment=tokenizer.decode(result["tokens"][1+len(left):1+len(left)+length]).strip() |
| if not fragment or fragment in seen or '\ufffd' in fragment or len(fragment.split())>14: |
| continue |
| seen.add(fragment) |
| rewritten=prefix+fragment+suffix |
| score=local_nll(model,tokenizer,rewritten,a,a+len(fragment),precision) |
| trace=[] |
| for row in result["trace"]: |
| |
| trace.append({"step":row["step"],"remaining":row["remaining"],"model_text":row["text"]}) |
| candidates.append({"replacement":fragment,"text":rewritten,"local_nll":score, |
| "delta_nll":score-before,"seed":seed+attempt,"masked_tokens":length,"trace":trace}) |
| candidates.sort(key=lambda c:c["local_nll"]) |
| result={"original":text,"selection":{"start":a,"end":b,"text":text[a:b]}, |
| "prefix":prefix,"suffix":suffix,"original_local_nll":before,"candidates":candidates, |
| "attempts":len(lengths),"latency_ms":(time.perf_counter()-started)*1000, |
| "method":"Bidirectional denoising; rank by causal NLL of the span and up to 16 following tokens.", |
| "limitation":"Ranking reflects this model's preferences, not verified grammar or semantic correctness."} |
| if ranker is not None: |
| ranker.rank(result) |
| else: |
| best=candidates[0] if candidates else None |
| replace=bool(best and best['local_nll']<before) |
| result['recommendation']={'action':'replace' if replace else 'keep', |
| 'replacement':best['replacement'] if replace else text[a:b],'reason':'local_likelihood'} |
| result['latency_ms']=(time.perf_counter()-started)*1000 |
| return result |
|
|
