lingbotworld2_t5 ¶

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply the tanh GELU approximation."""
    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))

def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs,
) -> BaseEncoderOutput:
    """Encode token ids and return source-compatible hidden states."""
    del position_ids, inputs_embeds, output_hidden_states, kwargs
    assert input_ids is not None
    x = self.dropout(self.token_embedding(input_ids))
    for block in self.blocks:
        x = block(x, attention_mask)
    x = self.dropout(self.norm(x))
    return BaseEncoderOutput(last_hidden_state=x, attention_mask=attention_mask)

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
    """Load source `.pth` weights whose names already match this module."""
    params_dict = dict(self.named_parameters())
    loaded_params: set[str] = set()
    for name, loaded_weight in weights:
        if name not in params_dict:
            continue
        param = params_dict[name]
        weight_loader = getattr(param, "weight_loader", default_weight_loader)
        weight_loader(param, loaded_weight)
        loaded_params.add(name)
    return loaded_params

def forward(
    self,
    x: torch.Tensor,
    context: torch.Tensor | None = None,
    mask: torch.Tensor | None = None,
    pos_bias: torch.Tensor | None = None,
) -> torch.Tensor:
    """Project QKV, add relative bias/mask, and return attended states."""
    context = x if context is None else context
    b, n, c = x.size(0), self.num_heads, self.head_dim
    q = self.q(x).view(b, -1, n, c)
    k = self.k(context).view(b, -1, n, c)
    v = self.v(context).view(b, -1, n, c)
    attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
    if pos_bias is not None:
        attn_bias += pos_bias
    if mask is not None:
        assert mask.ndim in (2, 3)
        mask = mask.view(b, 1, 1, -1) if mask.ndim == 2 else mask.unsqueeze(1)
        attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
    attn = torch.einsum("binc,bjnc->bnij", q, k) + attn_bias
    attn = F.softmax(attn.float(), dim=-1).type_as(attn)
    x = torch.einsum("bnij,bjnc->binc", attn, v)
    return self.dropout(self.o(x.reshape(b, -1, n * c)))

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply the gated feed-forward projection."""
    x = self.fc1(x) * self.gate(x)
    x = self.dropout(x)
    x = self.fc2(x)
    return self.dropout(x)

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Normalize in fp32 and apply the learned scale."""
    x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
    if self.weight.dtype in (torch.float16, torch.bfloat16):
        x = x.type_as(self.weight)
    return self.weight * x

def forward(self, lq: int, lk: int) -> torch.Tensor:
    """Build a relative-position bias tensor for attention logits."""
    device = self.embedding.weight.device
    rel_pos = torch.arange(lk, device=device).unsqueeze(0) - torch.arange(lq, device=device).unsqueeze(1)
    rel_pos = self._relative_position_bucket(rel_pos)
    rel_pos_embeds = self.embedding(rel_pos)
    return rel_pos_embeds.permute(2, 0, 1).unsqueeze(0).contiguous()

def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None, pos_bias: torch.Tensor | None = None) -> torch.Tensor:
    """Run self-attention and feed-forward residual updates."""
    e = pos_bias if self.pos_embedding is None else self.pos_embedding(x.size(1), x.size(1))
    x = self._fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
    return self._fp16_clamp(x + self.ffn(self.norm2(x)))