hidden_markov_part1.py

class HiddenMarkovModel:
    def __init__(
        self, environment: ConllEnvironment, smooth_factor: float = 1e-10
    ) -> None:
        self.Q = environment.states
        self.V = environment.observations
        self._reversed_V = environment.to_reverse_dict(self.V)
        # A = transition probability matrix
        self.A = np.zeros((len(self.Q), len(self.Q)))
        # B = probability of observation from state (p(o|s))
        self.B = np.zeros((len(self.Q), len(self.V)))
        # Pi = initial probability distribution
        self.Pi = np.zeros(len(self.Q))
        self.smooth_factor = smooth_factor

    def train(self, environment: ConllEnvironment) -> None:
        previous = -1
        for i, s, o in environment.transit():
            self.B[s][o] += 1
            if i == 0:
                self.Pi[s] += 1
            else:
                self.A[previous][s] += 1
            previous = s

        self.A = self.count_to_probability(self.A)
        self.B = self.count_to_probability(self.B)
        self.Pi = self.count_to_probability(self.Pi)

    def count_to_probability(self, count_matrix: NDArray) -> NDArray:

        smoothed_count = count_matrix + self.smooth_factor
        if len(count_matrix.shape) > 1:
            probability_matrix = smoothed_count / smoothed_count.sum(
                axis=1, keepdims=True
            )
        else:
            probability_matrix = smoothed_count / smoothed_count.sum()
        return probability_matrix