from train.train import train_model
from model.encoder import EncoderModel

n_tokens = 4
max_count = 3

model = EncoderModel(
    input_dim=n_tokens,
    hidden_dim=128,
    d_ff=256,
    output_dim=max_count + 1,
    n_layers=6,
    num_heads=8,
    # use_positional=False,
    # use_feedforward=False
)
train_model(
    model,
    lr=0.00001,
    num_steps=1000,
    batch_size=10,
    n_tokens=n_tokens,
    seqlen=16,
    max_count=max_count
)