Spaces:
Sleeping
Sleeping
| """ | |
| Sinusoidal Positional Encoding | |
| """ | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| class PositionalEncoding(nn.Module): | |
| """ | |
| Sinusoidal Positional Encoding | |
| PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) | |
| PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model)) | |
| """ | |
| def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1): | |
| """ | |
| Args: | |
| d_model: Model dimension | |
| max_len: Maximum sequence length | |
| dropout: Dropout rate | |
| """ | |
| super().__init__() | |
| self.dropout = nn.Dropout(dropout) if dropout > 0 else None | |
| # Create positional encoding matrix | |
| pe = torch.zeros(max_len, d_model) | |
| position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) | |
| # Compute the div_term: 10000^(2i/d_model) | |
| div_term = torch.exp( | |
| torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model) | |
| ) | |
| # Apply sine to even indices | |
| pe[:, 0::2] = torch.sin(position * div_term) | |
| # Apply cosine to odd indices | |
| if d_model % 2 == 0: | |
| pe[:, 1::2] = torch.cos(position * div_term) | |
| else: | |
| # Handle odd d_model | |
| pe[:, 1::2] = torch.cos(position * div_term[:-1]) | |
| # Add batch dimension: [1, max_len, d_model] | |
| pe = pe.unsqueeze(0) | |
| # Register as buffer (not a parameter, but part of state_dict) | |
| self.register_buffer('pe', pe) | |
| def forward(self, x): | |
| """ | |
| Args: | |
| x: [batch_size, seq_len, d_model] | |
| Returns: | |
| x with positional encoding added: [batch_size, seq_len, d_model] | |
| """ | |
| seq_len = x.size(1) | |
| x = x + self.pe[:, :seq_len, :] | |
| if self.dropout is not None: | |
| x = self.dropout(x) | |
| return x | |
| class LearnedPositionalEncoding(nn.Module): | |
| """ | |
| Learned Positional Encoding (alternative to sinusoidal) | |
| Can potentially learn better position representations for specific tasks | |
| """ | |
| def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1): | |
| """ | |
| Args: | |
| d_model: Model dimension | |
| max_len: Maximum sequence length | |
| dropout: Dropout rate | |
| """ | |
| super().__init__() | |
| self.dropout = nn.Dropout(dropout) if dropout > 0 else None | |
| self.pe = nn.Parameter(torch.randn(1, max_len, d_model) * 0.02) | |
| def forward(self, x): | |
| """ | |
| Args: | |
| x: [batch_size, seq_len, d_model] | |
| Returns: | |
| x with positional encoding added: [batch_size, seq_len, d_model] | |
| """ | |
| seq_len = x.size(1) | |
| x = x + self.pe[:, :seq_len, :] | |
| if self.dropout is not None: | |
| x = self.dropout(x) | |
| return x | |