# Buggy code
def __init__(
    self,
    hidden_size: int,
    embeddings: TokenEmbeddings,
    tag_dictionary: Dictionary,
    tag_type: str,
    use_crf: bool = True,
    use_rnn: bool = True,
    rnn_layers: int = 1,
    dropout: float = 0.0,
    word_dropout: float = 0.05,
    locked_dropout: float = 0.5,
    train_initial_hidden_state: bool = False,
    rnn_type: str = "LSTM",
    pickle_module: str = "pickle",
    beta: float = 1.0,
    loss_weights: Dict[str, float] = None,
):
    """
    Initializes a SequenceTagger
    :param hidden_size: number of hidden states in RNN
    :param embeddings: word embeddings used in tagger
    :param tag_dictionary: dictionary of tags you want to predict
    :param tag_type: string identifier for tag type
    :param use_crf: if True use CRF decoder, else project directly to tag space
    :param use_rnn: if True use RNN layer, otherwise use word embeddings directly
    :param rnn_layers: number of RNN layers
    :param dropout: dropout probability
    :param word_dropout: word dropout probability
    :param locked_dropout: locked dropout probability
    :param train_initial_hidden_state: if True, trains initial hidden state of RNN
    :param beta: Parameter for F-beta score for evaluation and training annealing
    :param loss_weights: Dictionary of weights for classes (tags) for the loss function
    (if any tag's weight is unspecified it will default to 1.0)

    """

    super(SequenceTagger, self).__init__()
    self.use_rnn = use_rnn
    self.hidden_size = hidden_size
    self.use_crf: bool = use_crf
    self.rnn_layers: int = rnn_layers

    self.trained_epochs: int = 0

    self.embeddings = embeddings

    # set the dictionaries
    self.tag_dictionary: Dictionary = tag_dictionary
    self.tag_type: str = tag_type
    self.tagset_size: int = len(tag_dictionary)

    self.beta = beta

    self.weight_dict = loss_weights
    # Initialize the weight tensor
    if loss_weights is not None:
        n_classes = len(self.tag_dictionary)
        weight_list = [1.0 for i in range(n_classes)]
        for i, tag in enumerate(self.tag_dictionary.get_items()):
            if tag in loss_weights.keys():
                weight_list[i] = loss_weights[tag]
        self.loss_weights = torch.FloatTensor(weight_list).to(flair.device)
    else:
        self.loss_weights = None

    # initialize the network architecture
    self.nlayers: int = rnn_layers
    self.hidden_word = None

    # dropouts
    self.use_dropout: float = dropout
    self.use_word_dropout: float = word_dropout
    self.use_locked_dropout: float = locked_dropout

    self.pickle_module = pickle_module

    if dropout > 0.0:
        self.dropout = torch.nn.Dropout(dropout)

    if word_dropout > 0.0:
        self.word_dropout = flair.nn.WordDropout(word_dropout)

    if locked_dropout > 0.0:
        self.locked_dropout = flair.nn.LockedDropout(locked_dropout)

    rnn_input_dim: int = self.embeddings.embedding_length

    self.relearn_embeddings: bool = True

    if self.relearn_embeddings:
        self.embedding2nn = torch.nn.Linear(rnn_input_dim, rnn_input_dim)

    self.train_initial_hidden_state = train_initial_hidden_state
    self.bidirectional = True
    self.rnn_type = rnn_type

    # bidirectional LSTM on top of embedding layer
    if self.use_rnn:
        num_directions = 2 if self.bidirectional else 1

        if self.rnn_type in ["LSTM", "GRU"]:
            self.rnn = getattr(torch.nn, self.rnn_type)(
                rnn_input_dim,
                hidden_size,
                num_layers=self.nlayers,
                dropout=0.0 if self.nlayers == 1 else 0.5,
                bidirectional=True,
                batch_first=True,
            )
            # Create initial hidden state and initialize it
            if self.train_initial_hidden_state:
                self.hs_initializer = torch.nn.init.xavier_normal_

                self.lstm_init_h = Parameter(
                    torch.randn(self.nlayers * num_directions, self.hidden_size),
                    requires_grad=True,
                )

                self.lstm_init_c = Parameter(
                    torch.randn(self.nlayers * num_directions, self.hidden_size),
                    requires_grad=True,
                )

                # TODO: Decide how to initialize the hidden state variables
                # self.hs_initializer(self.lstm_init_h)
                # self.hs_initializer(self.lstm_init_c)

        # final linear map to tag space
        self.linear = torch.nn.Linear(hidden_size * num_directions, len(tag_dictionary))
    else:
        self.linear = torch.nn.Linear(
            self.embeddings.embedding_length, len(tag_dictionary)
        )

    if self.use_crf:
        self.transitions = torch.nn.Parameter(
            torch.randn(self.tagset_size, self.tagset_size)
        )

        self.transitions.detach()[
            self.tag_dictionary.get_idx_for_item(START_TAG), :
        ] = -10000

        self.transitions.detach()[
            :, self.tag_dictionary.get_idx_for_item(STOP_TAG)
        ] = -10000

    self.to(flair.device)


# Traceback (AssertionError)
...
2020-04-30 14:36:40,370 epoch 1 - iter 96/120 - loss 209973.64290365 - samples/sec: 33.04
2020-04-30 14:36:52,323 epoch 1 - iter 108/120 - loss 208670.76244213 - samples/sec: 32.17
2020-04-30 14:37:03,579 epoch 1 - iter 120/120 - loss 206854.72747396 - samples/sec: 34.17
2020-04-30 14:37:03,594 ----------------------------------------------------------------------------------------------------
2020-04-30 14:37:03,594 EPOCH 1 done: loss 206854.7275 - lr 0.1000000
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-18-67e2f7ae9c06> in <module>()
15 trainer.train('./resource/taggers/swahili-dumb-ner',
16               learning_rate=0.1,
---> 17               mini_batch_size=32)
18
19 # Plotting the performance graphs

3 frames
/usr/local/lib/python3.6/dist-packages/flair/models/sequence_tagger_model.py in _viterbi_decode(self, feats, transitions, all_scores)
754
755         start = best_path.pop()
--> 756         assert start == id_start
757         best_path.reverse()
758

AssertionError:

# Fixed code
def __init__(
    self,
    hidden_size: int,
    embeddings: TokenEmbeddings,
    tag_dictionary: Dictionary,
    tag_type: str,
    use_crf: bool = True,
    use_rnn: bool = True,
    rnn_layers: int = 1,
    dropout: float = 0.0,
    word_dropout: float = 0.05,
    locked_dropout: float = 0.5,
    train_initial_hidden_state: bool = False,
    rnn_type: str = "LSTM",
    pickle_module: str = "pickle",
    beta: float = 1.0,
    loss_weights: Dict[str, float] = None,
):
    """
    Initializes a SequenceTagger
    :param hidden_size: number of hidden states in RNN
    :param embeddings: word embeddings used in tagger
    :param tag_dictionary: dictionary of tags you want to predict
    :param tag_type: string identifier for tag type
    :param use_crf: if True use CRF decoder, else project directly to tag space
    :param use_rnn: if True use RNN layer, otherwise use word embeddings directly
    :param rnn_layers: number of RNN layers
    :param dropout: dropout probability
    :param word_dropout: word dropout probability
    :param locked_dropout: locked dropout probability
    :param train_initial_hidden_state: if True, trains initial hidden state of RNN
    :param beta: Parameter for F-beta score for evaluation and training annealing
    :param loss_weights: Dictionary of weights for classes (tags) for the loss function
    (if any tag's weight is unspecified it will default to 1.0)

    """

    super(SequenceTagger, self).__init__()
    self.use_rnn = use_rnn
    self.hidden_size = hidden_size
    self.use_crf: bool = use_crf
    self.rnn_layers: int = rnn_layers

    self.trained_epochs: int = 0

    self.embeddings = embeddings

    # set the dictionaries
    self.tag_dictionary: Dictionary = tag_dictionary
    # if we use a CRF, we must add special START and STOP tags to the dictionary
    if use_crf:
        self.tag_dictionary.add_item(START_TAG)
        self.tag_dictionary.add_item(STOP_TAG)

    self.tag_type: str = tag_type
    self.tagset_size: int = len(tag_dictionary)

    self.beta = beta

    self.weight_dict = loss_weights
    # Initialize the weight tensor
    if loss_weights is not None:
        n_classes = len(self.tag_dictionary)
        weight_list = [1.0 for i in range(n_classes)]
        for i, tag in enumerate(self.tag_dictionary.get_items()):
            if tag in loss_weights.keys():
                weight_list[i] = loss_weights[tag]
        self.loss_weights = torch.FloatTensor(weight_list).to(flair.device)
    else:
        self.loss_weights = None

    # initialize the network architecture
    self.nlayers: int = rnn_layers
    self.hidden_word = None

    # dropouts
    self.use_dropout: float = dropout
    self.use_word_dropout: float = word_dropout
    self.use_locked_dropout: float = locked_dropout

    self.pickle_module = pickle_module

    if dropout > 0.0:
        self.dropout = torch.nn.Dropout(dropout)

    if word_dropout > 0.0:
        self.word_dropout = flair.nn.WordDropout(word_dropout)

    if locked_dropout > 0.0:
        self.locked_dropout = flair.nn.LockedDropout(locked_dropout)

    rnn_input_dim: int = self.embeddings.embedding_length

    self.relearn_embeddings: bool = True

    if self.relearn_embeddings:
        self.embedding2nn = torch.nn.Linear(rnn_input_dim, rnn_input_dim)

    self.train_initial_hidden_state = train_initial_hidden_state
    self.bidirectional = True
    self.rnn_type = rnn_type

    # bidirectional LSTM on top of embedding layer
    if self.use_rnn:
        num_directions = 2 if self.bidirectional else 1

        if self.rnn_type in ["LSTM", "GRU"]:
            self.rnn = getattr(torch.nn, self.rnn_type)(
                rnn_input_dim,
                hidden_size,
                num_layers=self.nlayers,
                dropout=0.0 if self.nlayers == 1 else 0.5,
                bidirectional=True,
                batch_first=True,
            )
            # Create initial hidden state and initialize it
            if self.train_initial_hidden_state:
                self.hs_initializer = torch.nn.init.xavier_normal_

                self.lstm_init_h = Parameter(
                    torch.randn(self.nlayers * num_directions, self.hidden_size),
                    requires_grad=True,
                )

                self.lstm_init_c = Parameter(
                    torch.randn(self.nlayers * num_directions, self.hidden_size),
                    requires_grad=True,
                )

                # TODO: Decide how to initialize the hidden state variables
                # self.hs_initializer(self.lstm_init_h)
                # self.hs_initializer(self.lstm_init_c)

        # final linear map to tag space
        self.linear = torch.nn.Linear(hidden_size * num_directions, len(tag_dictionary))
    else:
        self.linear = torch.nn.Linear(
            self.embeddings.embedding_length, len(tag_dictionary)
        )

    if self.use_crf:
        self.transitions = torch.nn.Parameter(
            torch.randn(self.tagset_size, self.tagset_size)
        )

        self.transitions.detach()[
            self.tag_dictionary.get_idx_for_item(START_TAG), :
        ] = -10000

        self.transitions.detach()[
            :, self.tag_dictionary.get_idx_for_item(STOP_TAG)
        ] = -10000

    self.to(flair.device)