File size: 3,640 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 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 | """Pretrained-text conditioning with zero-initialized residual cross-attention.
Frozen T5 features are inputs, not vocabulary IDs. Visual/action positions and
vocabulary stay unchanged. No-text requests bypass the added residual exactly.
"""
import torch
_UNSET = object()
from torch import nn
import torch.nn.functional as F
class TextResidual(nn.Module):
def __init__(self, hidden=768, text_width=768, inner=256, heads=4):
super().__init__();self.heads=heads;self.inner=inner
self.query_norm=nn.LayerNorm(hidden)
self.text_norm=nn.LayerNorm(text_width)
self.q=nn.Linear(hidden,inner,bias=False)
self.k=nn.Linear(text_width,inner,bias=False)
self.v=nn.Linear(text_width,inner,bias=False)
self.out=nn.Linear(inner,hidden,bias=False)
nn.init.zeros_(self.out.weight)
def forward(self, hidden, text, mask):
b,t,_=hidden.shape
if text.shape[0]!=b:
assert text.shape[0]==1, 'Generation expansion is valid only within a single episode'
text=text.expand(b,-1,-1);mask=mask.expand(b,-1)
dim=self.inner//self.heads
q=self.q(self.query_norm(hidden)).view(b,t,self.heads,dim).transpose(1,2)
memory=self.text_norm(text)
k=self.k(memory).view(b,-1,self.heads,dim).transpose(1,2)
v=self.v(memory).view(b,-1,self.heads,dim).transpose(1,2)
value=F.scaled_dot_product_attention(q,k,v,attn_mask=mask[:,None,None,:].bool(),dropout_p=0.)
value=value.transpose(1,2).reshape(b,t,self.inner)
# Explicitly suppress all-padding / deliberately dropped instruction.
value=self.out(value)*mask.any(-1)[:,None,None].to(value.dtype)
return hidden+value.to(hidden.dtype)
class TextWorldModel(nn.Module):
def __init__(self, core, text_width=768, layers=(3,7,11), inner=256):
super().__init__();self.core=core;self.layers=list(layers)
self.adapters=nn.ModuleDict({str(i):TextResidual(core.config.hidden_size,text_width,inner) for i in layers})
self._text=None;self._mask=None;self._handles=[]
for i in layers:
self._handles.append(core.model.layers[i].register_forward_pre_hook(self._hook(i),with_kwargs=True))
def _hook(self,index):
def hook(module,args,kwargs):
if self._text is None:return args,kwargs
adapter=self.adapters[str(index)]
if args:
return (adapter(args[0],self._text,self._mask),)+args[1:],kwargs
kw=dict(kwargs);kw['hidden_states']=adapter(kw['hidden_states'],self._text,self._mask)
return args,kw
return hook
def set_text(self,text,mask):
assert (text is None)==(mask is None)
if text is not None:
assert text.ndim==3 and mask.shape==text.shape[:2]
assert torch.isfinite(text).all()
self._text=text;self._mask=mask
def forward(self,text_features=_UNSET,text_mask=_UNSET,**kwargs):
# Generation and selected-token scoring use the explicitly bound episode.
# No activation checkpointing: mutable conditioning cannot change during replay.
if text_features is not _UNSET:
assert text_mask is not _UNSET
self.set_text(text_features,text_mask)
else:
assert text_mask is _UNSET
return self.core(**kwargs)
def generate(self,*args,**kwargs):
# Caller sets the same episode instruction before every generated frame.
return self.core.generate(*args,**kwargs)
def compute_transition_scores(self,*args,**kwargs):return self.core.compute_transition_scores(*args,**kwargs)
|