logits_penalize_part_crf

range_labels = tf.range(0, self.target_vocab_size)
res = lambda inp, val: tf.reduce_sum(tf.cast(tf.equal(inp, val), tf.float32))
count_arr = tf.convert_to_tensor([res(y, i) + 1 for i in range_labels])
priores = count_arr / tf.reduce_sum(count_arr)
priores = tf.clip_by_value(priores, -1e12, 1e12)
label_sequences = tf.convert_to_tensor(y, dtype=tf.int32)
logits = tf.nn.log_softmax(logits) - tf.math.log(priores)
log_likelihood, trans_params = tfa.text.crf.crf_log_likelihood(inputs=logits,
                                                                           tag_indices=label_sequences,
                                                                           sequence_lengths=seq_length)
self.transition_params = tf.Variable(self.transition_params, trainable=False)
loss = tf.reduce_mean(-log_likelihood)