def rnn_encoder(n_base, state_len, insize=1, first_conv_size=4, stride=5, winlen=19, activation='swish', rnn_type='lstm', features=768, scale=5.0, blank_score=None, expand_blanks=True, num_layers=5, norm=None):
    rnn = layers[rnn_type]
    return Serial([
        conv(insize, first_conv_size, ks=5, bias=True, activation=activation, norm=norm),
        conv(first_conv_size, 16, ks=5, bias=True, activation=activation, norm=norm),
        conv(16, features, ks=winlen, stride=stride, bias=True, activation=activation, norm=norm),
        Permute([2, 0, 1]),
        *(rnn(features, features, reverse=(num_layers - i) % 2) for i in range(num_layers)),
        LinearCRFEncoder(
            features, n_base, state_len, activation='tanh', scale=scale,
            blank_score=blank_score, expand_blanks=expand_blanks
        )
    ])

[global_norm]
# State Length of CRF, determine how many state CRF decoder need to consider
state_len = 4

[input]
features = 1


[labels]
labels = [ "N", "A", "C", "G", "T",]
# labels of bases, N means empty state.

[qscore]
scale = 0.9356
bias = -0.1721
# bias factor of Q Score

[model]
package = "bonito.crf"
# use CRF or CTC at decoder (since bonito v0.4, only CRF is valid)

[encoder]
scale = 5.0
rnn_type = "lstm"
winlen = 19
features = 384
activation = "swish"
stride = 5
# downsample stride
blank_score = 2.0

[basecaller]
# config when basecalling, don't affect training
batchsize = 512
chunksize = 10000
overlap = 500

@register
class Convolution(Module):

    def __init__(self, insize, size, winlen, stride=1, padding=0, bias=True, activation=None, norm=None):
        super().__init__()
        self.conv = torch.nn.Conv1d(insize, size, winlen, stride=stride, padding=padding, bias=bias)
        self.activation = layers.get(activation, lambda: activation)()
        if isinstance(norm, dict):
            self.norm = from_dict(norm)
        elif isinstance(norm, str):
            self.norm = layers[norm](size)
        else:
            self.norm = norm

    def forward(self, x):
        h = self.conv(x)
        if self.norm is not None:
            h = self.norm(h)
        if self.activation is not None:
            h = self.activation(h)
        return h

    def to_dict(self, include_weights=False):
        res = {
            "insize": self.conv.in_channels,
            "size": self.conv.out_channels,
            "bias": self.conv.bias is not None,
            "winlen": self.conv.kernel_size[0],
            "stride": self.conv.stride[0],
            "padding": self.conv.padding[0],
        }
        if self.activation is not None:
            res["activation"] = self.activation.name
        if self.norm is not None:
            res["norm"] = to_dict(self.norm, include_weights)
            #simplify default case e.g. norm="batchnorm"
            if not include_weights and self.norm.name in layers:
                if res["norm"] == to_dict(layers[self.norm.name](res["size"])):
                    res["norm"] = self.norm.name

        if include_weights:
            res['params'] = {
                'W': self.conv.weight, 'b': self.conv.bias if self.conv.bias is not None else []
            }
        return res

@register
class LSTM(RNNWrapper):

    def __init__(self, size, insize, bias=True, reverse=False):
        super().__init__(torch.nn.LSTM, size, insize, bias=bias, reverse=reverse)

    def to_dict(self, include_weights=False):
        res = {
            'size': self.rnn.hidden_size,
            'insize': self.rnn.input_size,
            'bias': self.rnn.bias,
            'reverse': self.reverse,
        }
        if include_weights:
            res['params'] = {
                'iW': self.rnn.weight_ih_l0.reshape(4, self.rnn.hidden_size, self.rnn.input_size),
                'sW': self.rnn.weight_hh_l0.reshape(4, self.rnn.hidden_size, self.rnn.hidden_size),
                'b': self.rnn.bias_ih_l0.reshape(4, self.rnn.hidden_size)
            }
        return res

@register
class LinearCRFEncoder(Module):

    def __init__(self, insize, n_base, state_len, bias=True, scale=None, activation=None, blank_score=None, expand_blanks=True, permute=None):
        super().__init__()
        self.scale = scale
        self.n_base = n_base
        self.state_len = state_len
        self.blank_score = blank_score
        self.expand_blanks = expand_blanks
        size = (n_base + 1) * n_base**state_len if blank_score is None else n_base**(state_len + 1)
        self.linear = torch.nn.Linear(insize, size, bias=bias)
        self.activation = layers.get(activation, lambda: activation)()
        self.permute = permute

    def forward(self, x):
        if self.permute is not None:
            x = x.permute(*self.permute)
        scores = self.linear(x)
        if self.activation is not None:
            scores = self.activation(scores)
        if self.scale is not None:
            scores = scores * self.scale
        if self.blank_score is not None and self.expand_blanks:
            T, N, C = scores.shape
            scores = torch.nn.functional.pad(
                scores.view(T, N, C // self.n_base, self.n_base),
                (1, 0, 0, 0, 0, 0, 0, 0),
                value=self.blank_score
            ).view(T, N, -1)
        return scores

def compute_scores(model, batch, beam_width=32, beam_cut=100.0, scale=1.0, offset=0.0, blank_score=2.0, reverse=False):
    """
    Compute scores for model.
    """
    with torch.inference_mode():
        device = next(model.parameters()).device
        dtype = torch.float16 if half_supported() else torch.float32
        scores = model(batch.to(dtype).to(device))
        if reverse:
            scores = model.seqdist.reverse_complement(scores)
        with torch.cuda.device(scores.device):
            sequence, qstring, moves = beam_search(
                scores, beam_width=beam_width, beam_cut=beam_cut,
                scale=scale, offset=offset, blank_score=blank_score
            )
        return {
            'moves': moves,
            'qstring': qstring,
            'sequence': sequence,
        }

def compute_scores(model, batch, reverse=False):
    with torch.no_grad():
        device = next(model.parameters()).device
        dtype = torch.float16 if half_supported() else torch.float32
        scores = model.encoder(batch.to(dtype).to(device))
        if reverse: scores = model.seqdist.reverse_complement(scores)
        betas = model.seqdist.backward_scores(scores.to(torch.float32))
        trans, init = model.seqdist.compute_transition_probs(scores, betas)
    return {
        'trans': trans.to(dtype).transpose(0, 1),
        'init': init.to(dtype).unsqueeze(1),
    }
    
def beam_search_duplex(seq1, path1, t1, b1, seq2, path2, t2, b2, alphabet='NACGT', beamsize=5, pad=40, T=0.01):
    env = build_envelope(t1.shape[0], seq1, path1, t2.shape[0], seq2, path2, padding=pad)
    return crf_beam_search_duplex(
        t1, b1, t2, b2,
        alphabet=alphabet,
        beam_size=beamsize,
        beam_cut_threshold=T,
        envelope=env,
    )

import seqdist.sparse
from seqdist.ctc_simple import logZ_cupy, viterbi_alignments
from seqdist.core import SequenceDist, Max, Log, semiring

def backward_scores(self, scores, S: semiring=Log):
    T, N, _ = scores.shape
    Ms = scores.reshape(T, N, -1, self.n_base + 1)
    beta_T = Ms.new_full((N, self.n_base**(self.state_len)), S.one)
    return seqdist.sparse.bwd_scores_cupy(Ms, self.idx, beta_T, S, K=1)

from fast_ctc_decode import crf_greedy_search

def compute_scores(model, batch, beam_width=32, beam_cut=100.0, scale=1.0, offset=0.0, blank_score=2.0, reverse=False):
    """
    Compute scores for model.
    """
    with torch.inference_mode():
        device = next(model.parameters()).device
        dtype = torch.float16 if half_supported() else torch.float32
        scores = model(batch.to(dtype).to(device))
        if reverse:
            scores = model.seqdist.reverse_complement(scores)

        # switch dim 1 & 2
        scores_pad = scores.permute(1, 0, 2)

        # pad score
        n_base = model.seqdist.n_base
        T, N, C = scores_pad.shape
        scores_pad = torch.nn.functional.pad(
            scores_pad.view(T, N, C // n_base, n_base),
            (1, 0, 0, 0, 0, 0, 0, 0),
            value=blank_score
        ).view(T, N, -1)

        betas = model.seqdist.backward_scores(scores_pad.to(torch.float32))
        trans, init = model.seqdist.compute_transition_probs(scores_pad, betas)

        trans = trans.to(torch.float32).transpose(0, 1)
        init = init.to(torch.float32).unsqueeze(1)

        # offload
        tracebacks = trans.cpu()
        init = init.cpu()

        seq_tensor = torch.zeros((N, T), dtype=torch.uint8, device='cpu')
        qstring_tensor = torch.zeros((N, T), dtype=torch.uint8, device='cpu')
        moves_tensor = torch.zeros((N, T), dtype=torch.uint8, device='cpu')

        for batch_idx in range(N):
            tracebacks_batch = tracebacks[batch_idx].numpy()  # (T, 256, 5)
            init_batch = init[batch_idx][0].numpy()  # (256,)

            # greedy decode, cef_beam_search dont output qstring
            seq_batch, path_batch = crf_greedy_search(
                network_output=tracebacks_batch,
                init_state=init_batch,
                alphabet="NACGT",
                qstring=True,
                qscale=1,
                qbias=1
            )

            # re-encode
            seq_batch_str = seq_batch[:len(seq_batch) // 2]
            qstring_batch_str = seq_batch[len(seq_batch) // 2:]

            seq_as_numbers = np.frombuffer(seq_batch_str.encode(), dtype=np.uint8).copy()
            qstring_as_numbers = np.frombuffer(qstring_batch_str.encode(), dtype=np.uint8).copy()

            seq_tensor[batch_idx, path_batch] = torch.from_numpy(seq_as_numbers[:len(path_batch)])
            qstring_tensor[batch_idx, path_batch] = torch.from_numpy(qstring_as_numbers[:len(path_batch)])
            moves_tensor[batch_idx, path_batch] = 1

        return {
            'qstring': qstring_tensor,
            'sequence': seq_tensor,
            'moves': moves_tensor,
        }

# bonito train
positional arguments:
  training_directory

optional arguments:
  -h, --help            show this help message and exit
  --config CONFIG
  --pretrained PRETRAINED
  --directory DIRECTORY
  --device DEVICE
  --lr LR
  --seed SEED
  --epochs EPOCHS
  --batch BATCH
  --chunks CHUNKS
  --valid-chunks VALID_CHUNKS
  --no-amp
  -f, --force
  --restore-optim
  --nondeterministic
  --save-optim-every SAVE_OPTIM_EVERY
  --grad-accum-split GRAD_ACCUM_SPLIT
  --quantile-grad-clip
  --num-workers NUM_WORKERS

    def decode_batch(self, x):
        scores = self.seqdist.posteriors(x.to(torch.float32)) + 1e-8
        tracebacks = self.seqdist.viterbi(scores.log()).to(torch.int16).T
        return [self.seqdist.path_to_str(x) for x in tracebacks.cpu().numpy()]

bonito train \
--config config.toml \ # 前面提到的模型config，chunksize等会从数据集里自动读取
--device cuda:0 \ # 显卡
--epochs 5 # epoch
--lr 5e-4 \ # 学习率
--batch 96 \ # batchsize
--pretrained [email protected] \ # 现有模型，可以从ont提供的模型选，或者传入其他模型所在的目录
--directory dataset_dir/  \ # chunks.npy等所在的文件夹
model_dir/