Collision-1B / model /blocks.py
collision-10M's picture
Release COLLISION-10M with official in-house NLP Engine and analytics
ddc12cc verified
Raw History Blame Contribute Delete
1.07 kB
import torch
import torch.nn as nn
from model.attention import CausalSelfAttention
class FeedForward(nn.Module):
def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
super().__init__()
self.net = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Linear(d_ff, d_model),
nn.Dropout(dropout)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.net(x)
class TransformerBlock(nn.Module):
def __init__(self, d_model: int, n_head: int, max_seq_len: int, d_ff: int, dropout: float = 0.1):
super().__init__()
self.ln1 = nn.LayerNorm(d_model)
self.attn = CausalSelfAttention(d_model, n_head, max_seq_len, dropout)
self.ln2 = nn.LayerNorm(d_model)
self.ffn = FeedForward(d_model, d_ff, dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Pre-LN structure with residual connections
x = x + self.attn(self.ln1(x))
x = x + self.ffn(self.ln2(x))
return x