35 lines
939 B
Python
35 lines
939 B
Python
import torch
|
|
from ..Utils import fixed_positional_encoding
|
|
|
|
# WIP FOR BATCHING
|
|
class NanoSocratesEmbedder(torch.nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
vocabulary_size: int,
|
|
embedding_size: int
|
|
) -> None:
|
|
super().__init__()
|
|
self.__embedder = torch.nn.Embedding(
|
|
vocabulary_size,
|
|
embedding_size
|
|
)
|
|
|
|
def forward(self, tokenized_sentence: list[list[int]]) -> torch.Tensor:
|
|
|
|
TOKENIZED_TENSOR = torch.tensor(tokenized_sentence)
|
|
|
|
computed_embeddings: torch.Tensor = self.__embedder(TOKENIZED_TENSOR)
|
|
|
|
_ ,SENTENCE_LENGHT, EMBEDDING_SIZE = computed_embeddings.shape # for batching
|
|
|
|
POSITIONAL_ENCODINGS = fixed_positional_encoding(
|
|
SENTENCE_LENGHT,
|
|
EMBEDDING_SIZE
|
|
)
|
|
|
|
computed_embeddings = computed_embeddings + POSITIONAL_ENCODINGS.unsqueeze(0) # for batching
|
|
return computed_embeddings
|
|
|
|
|