# -*- coding: utf-8 -*-
# coding=utf-8
# Copyright 2019 The SGNMT Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This module contains predictors for n-gram (Kneser-Ney) language
modeling. This is a ``UnboundedVocabularyPredictor`` as the vocabulary
size ngram models normally do not permit complete enumeration of the
posterior.
"""
from cam.sgnmt.predictors.core import UnboundedVocabularyPredictor
from cam.sgnmt import utils
import math
try:
# Requires kenlm
import kenlm
except ImportError:
pass # Deal with it in decode.py
[docs]class KenLMPredictor(UnboundedVocabularyPredictor):
"""KenLM predictor based on
https://github.com/kpu/kenlm
The predictor state is described by the n-gram history.
"""
def __init__(self, path):
"""Creates a new n-gram language model predictor.
Args:
path (string): Path to the ARPA language model file
Raises:
NameError. If KenLM is not installed
"""
super(KenLMPredictor, self).__init__()
self.lm = kenlm.Model(path)
self.lm_state2 = kenlm.State()
[docs] def initialize(self, src_sentence):
"""Initializes the KenLM state.
Args:
src_sentence (list): Not used
"""
self.history = []
self._update_lm_state()
def _update_lm_state(self):
self.lm_state = kenlm.State()
tmp_state = kenlm.State()
self.lm.BeginSentenceWrite(self.lm_state)
for w in self.history[-6:]:
self.lm.BaseScore(self.lm_state, w, tmp_state)
self.lm_state, tmp_state = tmp_state, self.lm_state
[docs] def predict_next(self, words):
return {w: self.lm.BaseScore(self.lm_state,
"</s>" if w == utils.EOS_ID else str(w),
self.lm_state2)
for w in words}
[docs] def get_unk_probability(self, posterior):
"""Use the probability for '<unk>' in the language model """
return self.lm.BaseScore(self.lm_state, "<unk>", self.lm_state2)
[docs] def consume(self, word):
self.lm.BaseScore(self.lm_state, str(word), self.lm_state2)
self.lm_state, self.lm_state2 = self.lm_state2, self.lm_state
self.history.append(str(word))
def get_state(self):
return self.lm_state.clone()
[docs] def get_state(self):
"""Returns the current n-gram history """
return self.history
[docs] def set_state(self, state):
"""Sets the current n-gram history and LM state """
self.history = state
self._update_lm_state()
[docs] def is_equal(self, state1, state2):
return state1 == state2