Added a way to detach models and create them standalone

This commit is contained in:
Christian Risi
2025-10-10 18:43:20 +02:00
parent 15f203cad5
commit 92ae40013d
10 changed files with 164 additions and 20 deletions

View File

@@ -0,0 +1,32 @@
import torch
import Project_Model.Libs.Embedder as Embedder
from ..Classes import DeToken
class NanoSocraDecoder(torch.nn.Module):
def __init__(
self,
decoder_embedder: Embedder.NanoSocratesEmbedder,
decoder_layers: torch.nn.Sequential,
detokener: DeToken
) -> None:
super().__init__()
self.__decoder_embedder = decoder_embedder
self.__decoder = decoder_layers
self.__detokener = detokener
def forward(self, args: tuple[torch.Tensor, torch.Tensor]):
decoder_embedder_input, tgt_padding = args
decoder_tensor = self.__decoder_embedder(decoder_embedder_input)
decoder_output, _, _, _, _, _ = self.__decoder(
(decoder_tensor, decoder_tensor, decoder_tensor, tgt_padding, tgt_padding, False)
)
logits: torch.Tensor = self.__detokener(decoder_output)
return logits