| """Fixed architecture. Imports MLX only when the head is constructed.""" |
| def create_head(cfg): |
| import mlx.core as mx |
| import mlx.nn as nn |
| class JointSchemaHead(nn.Module): |
| def __init__(self): |
| super().__init__();d=cfg['head_dims'] |
| self.input_norm=nn.LayerNorm(cfg['hidden_size']);self.project=nn.Linear(cfg['hidden_size'],d) |
| self.evidence_query_norm=nn.LayerNorm(d);self.evidence_memory_norm=nn.LayerNorm(d) |
| self.evidence_attention=nn.MultiHeadAttention(d,cfg['attention_heads']) |
| self.joint=nn.TransformerEncoder(cfg['joint_layers'],d,cfg['attention_heads'],mlp_dims=cfg['ffn_dims'],dropout=0.0,activation=nn.gelu,norm_first=True) |
| self.score=nn.Linear(d,1) |
| def __call__(self,candidates,evidence,candidate_valid,evidence_valid): |
| q=self.project(self.input_norm(candidates.astype(mx.float32))) |
| m=self.project(self.input_norm(evidence.astype(mx.float32))) |
| |
| evidence_mask=mx.where(evidence_valid[:,None,None,:],mx.array(0.,mx.float32),mx.array(-1e9,mx.float32)) |
| candidate_mask=mx.where(candidate_valid[:,None,None,:],mx.array(0.,mx.float32),mx.array(-1e9,mx.float32)) |
| qn=self.evidence_query_norm(q);mn=self.evidence_memory_norm(m) |
| q=q+self.evidence_attention(qn,mn,mn,evidence_mask) |
| q=self.joint(q,mask=candidate_mask) |
| logits=self.score(q).squeeze(-1) |
| return mx.where(candidate_valid,logits,mx.array(-1e9,mx.float32)) |
| return JointSchemaHead() |
|
|
| def pack_numpy(features): |
| import numpy as np |
| if not features:raise ValueError('Empty batch') |
| hidden=features[0]['candidates'].shape[-1];n=max(len(f['candidates']) for f in features);m=max(len(f['evidence']) for f in features) |
| q=np.zeros((len(features),n,hidden),dtype=np.float32);e=np.zeros((len(features),m,hidden),dtype=np.float32) |
| qv=np.zeros((len(features),n),dtype=bool);ev=np.zeros((len(features),m),dtype=bool);labels=[] |
| for i,f in enumerate(features): |
| a,b=f['candidates'],f['evidence'];label=f['label'] |
| if a.ndim!=2 or b.ndim!=2 or len(a)<2 or len(b)<1 or a.shape[1]!=hidden or b.shape[1]!=hidden or not 0<=label<len(a):raise ValueError('Invalid feature shape/label') |
| if not np.isfinite(a).all() or not np.isfinite(b).all():raise ValueError('Nonfinite features') |
| q[i,:len(a)]=a;e[i,:len(b)]=b;qv[i,:len(a)]=True;ev[i,:len(b)]=True;labels.append(label) |
| return q,e,qv,ev,np.array(labels,dtype=np.int32) |
|
|
| def pack(features): |
| import mlx.core as mx |
| return tuple(mx.array(a) for a in pack_numpy(features)) |
|
|
