File size: 2,102 Bytes
64b9802 | 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 | """Causal ReCAP prompts shared by rollout and every policy log-probability pass."""
import torch
def retained(target,recent=6):
assert target>=2 and recent>=1
return [1]+list(range(max(2,target-recent),target))
def make_prefix(context,history,actions,target,recent=6):
kept=retained(target,recent)
assert max(kept)==target-1 and len(kept)==len(set(kept))
blocks=[context]
for t in kept:blocks += [history[t],actions[t].expand(context.shape[0],-1)]
return torch.cat(blocks,dim=1),kept
@torch.no_grad()
def rollout(model,context,anchor,actions,horizon,group,generator,visual=4375,width=80,temperature=1.0):
model.eval()
context=context.expand(group,-1);history={1:anchor.expand(group,-1)}
frames=[];traces=[]
for target in range(2,horizon+2):
prefix,kept=make_prefix(context,history,actions,target)
pending=prefix;kv=None;codes=[];lps=[]
for token in range(width):
out=model(input_ids=pending,past_key_values=kv,use_cache=True)
kv=out.past_key_values
logp=torch.log_softmax(out.logits[:,-1,:visual].float()/temperature,-1)
pending=torch.multinomial(logp.exp(),1,generator=generator)
codes.append(pending);lps.append(logp.gather(1,pending).squeeze(1))
response=torch.cat(codes,1)
frames.append({'prefix':prefix.detach(),'response':response.detach(),'sample_logp':torch.stack(lps,1).detach(),'target':target})
history[target]=response
traces.append({'target':target,'retained':kept,'prefix_length':prefix.shape[1],'positions_reset':True,'kv_reset':True,'future_observations':False})
return frames,traces
def frame_logp(model,prefix,response,visual=4375,temperature=1.0):
# Each target has its own rebased prefix. No discarded history is hidden in a cache.
ids=torch.cat([prefix,response[:,:-1]],1)
out=model(input_ids=ids,use_cache=False)
p=prefix.shape[1]
logits=out.logits[:,p-1:p+response.shape[1]-1,:visual].float()/temperature
return torch.log_softmax(logits,-1).gather(-1,response.unsqueeze(-1)).squeeze(-1)
|