aki.py

from torch import nn, Tensor
from huggingface_hub import PyTorchModelHubMixin
from typing import Optional


class AKIClassifier(nn.Module, PyTorchModelHubMixin):

    def __init__(self, input_size: int, n_classes: int):
        super(AKIClassifier, self).__init__()

        is_binary = n_classes <= 2

        layers = [
            self.nn_block(input_size, 256),
            self.nn_block(256),
            self.nn_block(256),
            self.nn_block(256),
            self.nn_block(256),
            self.nn_block(256),
            self.nn_block(256),
        ]

        if is_binary:
            layers.extend((nn.Linear(256, 1), nn.Sigmoid(), nn.Flatten(0, -1)))
        else:
            layers.extend((nn.Linear(256, n_classes), nn.Softmax(dim=-1)))
        self.layers = nn.Sequential(*layers)

    def nn_block(
        self, in_c: int, out_c: Optional[int] = None, dropout: float = 0.2
    ) -> nn.Module:
        return nn.Sequential(
            nn.Linear(in_c, out_c or in_c),
            nn.ReLU(),
            nn.Dropout(dropout),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.layers(x)