blob: bf5149e075b6ae4f8bee2229a63354dd09a74f51 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
|
from torch import Tensor, nn
from .l2_norm import l2_norm
class TokenEmbedding(nn.Module):
def __init__(self, num_tokens: int, dim: int, use_l2: bool = True) -> None:
super().__init__()
self.use_l2 = use_l2
self.to_embedding = nn.Embedding(num_tokens, dim)
if self.use_l2:
nn.init.normal_(self.to_embedding.weight, std=1e-5)
else:
nn.init.kaiming_normal_(self.to_embedding.weight)
def forward(self, x: Tensor) -> Tensor:
embedding = self.to_embedding(x)
return l2_norm(embedding) if self.use_l2 else embedding
|
