feat: Initial Commit

model/positional_encoding.py

@@ -0,0 +1,55 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+class PositionalEncoding(nn.Module):
+    """
+    Standard sinusoidal positional encoding from
+    "Attention Is All You Need".
+    """
+
+    def __init__(self, embed_dim, dropout=0.1, max_len=10000):
+        super().__init__()
+
+        self.dropout = nn.Dropout(dropout)
+
+        pe = torch.zeros(max_len, embed_dim)
+
+        position = torch.arange(
+            0,
+            max_len,
+            dtype=torch.float
+        ).unsqueeze(1)
+
+        div_term = torch.exp(
+            torch.arange(
+                0,
+                embed_dim,
+                2
+            ).float() *
+            (-math.log(10000.0) / embed_dim)
+        )
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        pe = pe.unsqueeze(0)
+
+        # Stored as a buffer so it moves with the model
+        # but isn't trained.
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        """
+        Args:
+            x: Tensor of shape
+               (batch_size, sequence_length, embed_dim)
+        """
+
+        seq_len = x.size(1)
+
+        x = x + self.pe[:, :seq_len]
+
+        return self.dropout(x)