| import torch | |
| import torch.nn as nn | |
| from huggingface_hub import PyTorchModelHubMixin | |
| class PodosTransformer(nn.Module,PyTorchModelHubMixin): | |
| def __init__(self, input_dim, model_dim, num_classes, num_heads=4, num_layers=2, dropout=0.1,temperature=1): | |
| super(PodosTransformer, self).__init__() | |
| self.temperature = temperature | |
| self.projection = nn.Linear(input_dim, model_dim) | |
| encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=num_heads, dropout=dropout) | |
| self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) | |
| self.fc = nn.Linear(model_dim, num_classes) | |
| def forward(self, x): | |
| x = self.projection(x) | |
| x = x.unsqueeze(1) | |
| x = self.transformer_encoder(x) | |
| x = x.mean(dim=1) | |
| x = self.fc(x) | |
| if self.temperature != 1.0: | |
| x = x / self.temperature | |
| return x |