| """Batched, answer-blind candidate selection extracted from measured code.""" |
| import ast |
| from collections import Counter |
| import re |
| import time |
| import numpy as np |
| import torch |
| import diffusion_runtime as base |
|
|
| def forbidden_tokens(context, window, ngram, max_count, eos): |
| """Constraints only inspect resolved left context; EOS is always permitted.""" |
| context = context[-window:] if window else [] |
| banned = {token for token, count in Counter(context).items() if max_count and count >= max_count} |
| if ngram and len(context) >= ngram - 1: |
| suffix = tuple(context[-(ngram - 1):]) if ngram > 1 else () |
| for i in range(len(context) - ngram + 1): |
| if tuple(context[i:i + ngram - 1]) == suffix: |
| banned.add(context[i + ngram - 1]) |
| return banned - set(eos) |
|
|
|
|
| @torch.inference_mode() |
| def sample_pool(sampler, prompt, cfg, length, seed): |
| """Batched candidate generation, with CPU rejection resampling included in timing.""" |
| model = sampler.model |
| eos_raw = model.config.eos_token_id |
| eos = set(eos_raw if isinstance(eos_raw,list) else [eos_raw]) |
| eos_one = min(eos) |
| pad = model.config.pad_token_id |
| p = prompt.shape[1] |
| count = cfg["best_of"] |
| wave_size = min(count, max(1,4096//(p+length))) |
| rng = np.random.default_rng(seed) |
| prefix_list = prompt[0].tolist() |
| outputs, forwards = [], 0 |
| for start in range(0,count,wave_size): |
| base.check_time(30) |
| batch = min(wave_size,count-start) |
| prefix = prompt.expand(batch,-1) |
| canvas = torch.full((batch,length),pad,device="cuda",dtype=torch.long) |
| mask = torch.ones_like(canvas,dtype=torch.bool) |
| cpu_ids = np.full((batch,length),-1,dtype=np.int64) |
| cpu_lp = np.zeros((batch,length),dtype=np.float64) |
| steps = min(cfg["steps"],length) |
| for step in range(steps): |
| base.check_time(30) |
| ids = torch.cat((prefix,canvas),1) |
| fullmask = torch.cat((torch.zeros_like(prefix,dtype=torch.bool),mask),1) |
| h = sampler.reference(ids,fullmask,p)[:,p:][mask] |
| logits = model.lm_head(h).float() |
| logits[:,pad] = -torch.inf |
| k = min(cfg["top_k"],logits.shape[-1]-1) |
| values, tokens = logits.topk(k,dim=-1) |
| |
| eos_values = logits[:,eos_one:eos_one+1] |
| eos_tokens = torch.full_like(tokens[:,:1],eos_one) |
| duplicate = tokens.eq(eos_one).any(-1,keepdim=True) |
| eos_values = eos_values.masked_fill(duplicate,-torch.inf) |
| values = torch.cat((values,eos_values),-1) |
| tokens = torch.cat((tokens,eos_tokens),-1) |
| normalizer = torch.logsumexp(logits/.9,-1).cpu().numpy() |
| raw_values = values.cpu().numpy() |
| token_values = tokens.cpu().numpy() |
| del logits, values, tokens, h |
| forwards += 1 |
| remaining_after = length*(steps-step-1)//steps |
| row_offset = 0 |
| for row in range(batch): |
| positions = np.flatnonzero(cpu_ids[row]<0) |
| commit = len(positions)-remaining_after |
| |
| maximum = raw_values[row_offset:row_offset+len(positions),0]/.9-normalizer[row_offset:row_offset+len(positions)] |
| chosen = np.argsort(maximum)[-commit:] |
| for local in sorted(chosen.tolist(),key=lambda x:int(positions[x])): |
| pos = int(positions[local]) |
| idx = row_offset+local |
| left = pos-1 |
| while left>=0 and cpu_ids[row,left]>=0: |
| left -= 1 |
| context = cpu_ids[row,left+1:pos].tolist() |
| if left<0: |
| context = prefix_list+context |
| window_context = context[-cfg["window"]:] if cfg["window"] else [] |
| banned = forbidden_tokens(context,cfg["window"],cfg["ngram"],cfg["max_repeat"],eos) |
| scores = raw_values[idx].astype(np.float64).copy() |
| if cfg["penalty"] != 1: |
| counts = Counter(window_context) |
| for j,token in enumerate(token_values[idx]): |
| if int(token) in counts and int(token) not in eos: |
| penalty = cfg["penalty"]**min(counts[int(token)],3) |
| scores[j] = scores[j]/penalty if scores[j]>=0 else scores[j]*penalty |
| scores /= .9 |
| for j,token in enumerate(token_values[idx]): |
| if int(token) in banned: |
| scores[j] = -np.inf |
| assert np.isfinite(scores).any(), "EOS must remain selectable" |
| probability = np.exp(scores-np.max(scores)) |
| probability /= probability.sum() |
| choice = int(rng.choice(len(scores),p=probability)) |
| token = int(token_values[idx,choice]) |
| assert token not in banned |
| cpu_ids[row,pos] = token |
| cpu_lp[row,pos] = float(raw_values[idx,choice]/.9-normalizer[idx]) |
| row_offset += len(positions) |
| known = cpu_ids>=0 |
| canvas = torch.as_tensor(np.where(known,cpu_ids,pad),device="cuda",dtype=torch.long) |
| mask = torch.as_tensor(~known,device="cuda") |
| for row in range(batch): |
| sequence = cpu_ids[row].tolist() |
| stop = next((i for i,token in enumerate(sequence) if token in eos),len(sequence)) |
| assert all(t>=0 for t in sequence[:stop]) |
| outputs.append({"ids":sequence[:stop],"ended":stop<len(sequence), |
| "mean_logprob":float(cpu_lp[row,:max(1,stop)].mean())}) |
| return outputs,forwards,wave_size |
|
|
|
|
| def code_candidate(text, repair): |
| blocks = re.findall(r"```(?:python)?\s*\n(.*?)```",text,flags=re.S) |
| source = blocks[0] if blocks else text.strip() |
| choices = [source] |
| if repair: |
| lines = source.splitlines() |
| choices += ["\n".join(lines[:i]) for i in range(len(lines)-1,max(0,len(lines)-17),-1)] |
| for candidate in choices: |
| try: |
| tree = ast.parse(candidate) |
| definitions = [n for n in tree.body if isinstance(n,ast.FunctionDef)] |
| forbidden = any(isinstance(n,(ast.Import,ast.ImportFrom,ast.Attribute,ast.ClassDef)) for n in ast.walk(tree)) |
| if definitions and not forbidden: |
| return candidate, [n.name for n in definitions], candidate!=source |
| except (SyntaxError,ValueError): |
| pass |
| return source,[],False |
|
|
|
|
| def rank_candidate(record, tokenizer, prompt_text, repair): |
| """No answer or tests argument: selector cannot access evaluation truth.""" |
| text = tokenizer.decode(record["ids"],skip_special_tokens=True).strip() |
| match = re.search(r"defining\s+([A-Za-z_]\w*)",prompt_text) |
| score = record["mean_logprob"] |
| changed = False |
| if match: |
| text,names,changed = code_candidate(text,repair) |
| score += 100 if match.group(1) in names else -100 |
| score += 10 if names else 0 |
| elif "only" in prompt_text.lower(): |
| words = text.split() |
| score += 20 if 1<=len(words)<=4 and len(text)<60 else -20 |
| health = base.health(record["ids"]) |
| score -= max(0,health["max_run"]-3)*.5 |
| score -= (1-health["distinct4"])*5 |
| if not text: |
| score -= 1000 |
| return score,text,changed |
|
|
|
|
| def generate_selected(sampler, tokenizer, prompt_text, cfg, seed, long=False): |
| initial = base.encode(tokenizer,prompt_text) |
| torch.cuda.synchronize() |
| start = time.perf_counter() |
| generated, gross, calls, chunks, repaired, max_batch = [],0,0,0,False,0 |
| selected_text = "" |
| target = max(1024,cfg["block"]) if long else min(256,cfg["block"]) |
| while len(generated)<target: |
| n = min(cfg["block"],target-len(generated)) |
| history = initial.new_tensor([generated[-256:]]) |
| prompt = torch.cat((initial,history),1) |
| pool,forwards,wave = sample_pool(sampler,prompt,cfg,n,seed+chunks*1009) |
| ranked = [(rank_candidate(row,tokenizer,prompt_text,cfg["repair"] and not long),index) for index,row in enumerate(pool)] |
| ranking,index = max(ranked,key=lambda item:item[0][0]) |
| selected = pool[index] |
| selected_text,changed = ranking[1],ranking[2] |
| repaired |= changed |
| gross += sum(len(row["ids"]) for row in pool) |
| calls += forwards |
| max_batch = max(max_batch,wave) |
| generated.extend(selected["ids"]) |
| chunks += 1 |
| if not long or selected["ended"]: |
| break |
| if long: |
| selected_text = tokenizer.decode(generated,skip_special_tokens=True) |
| delivered = len(generated) |
| else: |
| |
| delivered = len(tokenizer(selected_text,add_special_tokens=False)["input_ids"]) |
| torch.cuda.synchronize() |
| elapsed = time.perf_counter()-start |
| return selected_text,{"selected_tokens":delivered,"gross_candidate_tokens":gross,"seconds":elapsed, |
| "selected_tokens_per_second":delivered/elapsed,"gross_tokens_per_second":gross/elapsed, |
| "best_of":cfg["best_of"],"max_parallel_candidates":max_batch,"forward_calls":calls, |
| "chunks":chunks,"syntax_repaired":repaired,**base.health(generated)} |
|
|
|
|
| def cfg(block,steps=1,best_of=1,window=0,ngram=0,max_repeat=0,penalty=1.,top_k=50,repair=False): |
| return dict(block=block,steps=steps,best_of=best_of,window=window,ngram=ngram,max_repeat=max_repeat, |
| penalty=penalty,top_k=top_k,repair=repair,temperature=.9) |
|
|
|
|
