from .CTMeth_neural import Tokenizer
from .CTMeth_neural import TransformerClassifier
import torch.nn as nn

## This is the neural network CTMeth class used for training.
## It is a simple wrapper around the Tokenizer and TransformerClassifier classes.
## This approach is used to make the code more modular and easier to understand and gives possibility for future improvements.
## Evaluation module does not use this class, it uses the Tokenizer and TransformerClassifier classes directly.


class CTMeth_class(nn.Module):
    def __init__(self, params, *args, **kwargs):
        super(CTMeth_class, self).__init__()
        self.tokenizer = Tokenizer(**params)
        self.classifier = TransformerClassifier(
            sequence_length=self.tokenizer.sequence_length(**params), **params
        )

    def forward(self, x):
        x = self.tokenizer(x)
        return self.classifier(x)
