# Path: px/nmt/utils/trie_decoder_utils.py
# class TrieSamplerState(
#     collections.namedtuple('TrieSamplerState',
#                            ('cell_state', 'trie_keys', 'trie_exclude'))):
# class TrieSamplerAttentionState(
#     collections.namedtuple(
#         'TrieSamplerAttentionState',
#         ('cell_state', 'attention', 'time', 'alignments', 'alignment_history',
#          'attention_state', 'trie_keys', 'trie_exclude')),
#     attention_wrapper.AttentionWrapperState):
# class TrieBeamSearchDecoderState(
#     collections.namedtuple('TrieBeamSearchDecoderState',
#                            ('cell_state', 'log_probs', 'finished', 'lengths',
#                             'trie_keys', 'trie_exclude'))):
# class DecoderTrie(pygtrie.StringTrie):
# class TrieSamplerDecoder(basic_decoder.BasicDecoder):
# class TrieBeamSearchDecoder(beam_search_decoder.BeamSearchDecoder):
#   def __init__(self,
#                vocab_path,
#                eos_token='</s>',
#                unk_token='<unk>',
#                subword_option='',
#                subword_model=None,
#                prefix='',
#                optimize_ngrams_len=0):
#   def _create_vocab(self):
#   def _add_to_start_words(self, key):
#   def _create_start_words(self):
#   def _create_sentencepiece_processor(self):
#   def __getstate__(self):
#   def __setstate__(self, state):
#   def insert_question(self, question, answers):
#   def save_to_file(self, file_path):
#   def load_from_file(cls, file_path):
#   def populate_from_text_file(self, text_file_path):
# def _is_attention_state(state):
# def _is_gnmt_state(state):
#   def __init__(self, trie, trie_exclude, *args, **kwargs):
#   def initialize(self, name=None):
#   def step(self, time, inputs, state, name=None):
#   def __init__(self, trie, trie_exclude, *args, **kwargs):
#   def initialize(self, name=None):
#   def step(self, time, inputs, state, name=None):
# def _trie_beam_search_step(time, logits, next_cell_state, beam_state,
#                            batch_size, beam_width, end_token,
#                            length_penalty_weight, trie):
# def _init_trie_keys_py_func(beam_search=False):
#     def _py_func(batch_size, beam_size):
#     def _py_func(batch_size):
# def _amend_trie_keys_py_func(beam_search=False):
#     def _py_func(trie_keys, next_word_ids, next_beam_ids):
#     def _py_func(trie_keys, next_word_ids):
# def _clean_trie_key(trie, trie_key):
# def _trie_scores_py_func(trie, beam_search=False):
#     def _py_func(log_probs, sequence_lengths, scores, trie_keys, trie_exclude):
#     def _py_func(logits, trie_keys, trie_exclude):
#       def _py_func_inner(batch_idx):
# def _get_valid_continuation_idxs(trie, trie_exclude_batch_set, trie_key):
# def _create_trie_exclude_batch_set(trie, trie_exclude_batch):
# def _get_scores(log_probs, sequence_lengths, length_penalty_weight):
# def _length_penalty(sequence_lengths, penalty_factor):
# def _get_trie_scores(log_probs, sequence_lengths, length_penalty_weight,
#                      trie_keys, trie_exclude, trie):
# def _mask_probs(probs, eos_token, finished):
# def _check_maybe(t):
# def _maybe_tensor_gather_helper(gather_indices, gather_from, batch_size,
#                                 range_size, gather_shape):
# def _tensor_gather_helper(gather_indices,
#                           gather_from,
#                           batch_size,
#                           range_size,
#                           gather_shape,
#                           name=None):


import os.path
import shutil
import tempfile
import tensorflow as tf
from px.nmt.utils import trie_decoder_utils
from px.proto import aqa_pb2
        ('a b c b a c', 'y'),
        ('a c c b', 'z'),
    ]

    with open(cls.vocab_path, 'w') as f:
      f.write('\n'.join(cls.test_vocab))


  @classmethod
  def tearDownClass(cls):
    shutil.rmtree(cls.tmpdir)


  def testSaveLoadTrie(self):
    trie = trie_decoder_utils.DecoderTrie(TrieDecoderUtilsTest.vocab_path)
    for question, answer in TrieDecoderUtilsTest.test_trie_entries:
      trie.insert_question(question, answer)
    NEW_KEY = 'newkey'
    self.assertNotIn(NEW_KEY, trie)
    self.assertFalse(os.path.exists(TrieDecoderUtilsTest.save_path))
    trie.save_to_file(TrieDecoderUtilsTest.save_path)
    self.assertTrue(os.path.exists(TrieDecoderUtilsTest.save_path))
    trie[NEW_KEY] = 'foo'
    self.assertIn(NEW_KEY, trie)
    loaded_trie = trie_decoder_utils.DecoderTrie.load_from_file(
        TrieDecoderUtilsTest.save_path)
    self.assertNotIn(NEW_KEY, loaded_trie)

  def testCreateEmptyTrie(self):
    eos_token = '</s>'

# Path: px/nmt/utils/trie_decoder_utils.py
# class TrieSamplerState(
#     collections.namedtuple('TrieSamplerState',
#                            ('cell_state', 'trie_keys', 'trie_exclude'))):
# class TrieSamplerAttentionState(
#     collections.namedtuple(
#         'TrieSamplerAttentionState',
#         ('cell_state', 'attention', 'time', 'alignments', 'alignment_history',
#          'attention_state', 'trie_keys', 'trie_exclude')),
#     attention_wrapper.AttentionWrapperState):
# class TrieBeamSearchDecoderState(
#     collections.namedtuple('TrieBeamSearchDecoderState',
#                            ('cell_state', 'log_probs', 'finished', 'lengths',
#                             'trie_keys', 'trie_exclude'))):
# class DecoderTrie(pygtrie.StringTrie):
# class TrieSamplerDecoder(basic_decoder.BasicDecoder):
# class TrieBeamSearchDecoder(beam_search_decoder.BeamSearchDecoder):
#   def __init__(self,
#                vocab_path,
#                eos_token='</s>',
#                unk_token='<unk>',
#                subword_option='',
#                subword_model=None,
#                prefix='',
#                optimize_ngrams_len=0):
#   def _create_vocab(self):
#   def _add_to_start_words(self, key):
#   def _create_start_words(self):
#   def _create_sentencepiece_processor(self):
#   def __getstate__(self):
#   def __setstate__(self, state):
#   def insert_question(self, question, answers):
#   def save_to_file(self, file_path):
#   def load_from_file(cls, file_path):
#   def populate_from_text_file(self, text_file_path):
# def _is_attention_state(state):
# def _is_gnmt_state(state):
#   def __init__(self, trie, trie_exclude, *args, **kwargs):
#   def initialize(self, name=None):
#   def step(self, time, inputs, state, name=None):
#   def __init__(self, trie, trie_exclude, *args, **kwargs):
#   def initialize(self, name=None):
#   def step(self, time, inputs, state, name=None):
# def _trie_beam_search_step(time, logits, next_cell_state, beam_state,
#                            batch_size, beam_width, end_token,
#                            length_penalty_weight, trie):
# def _init_trie_keys_py_func(beam_search=False):
#     def _py_func(batch_size, beam_size):
#     def _py_func(batch_size):
# def _amend_trie_keys_py_func(beam_search=False):
#     def _py_func(trie_keys, next_word_ids, next_beam_ids):
#     def _py_func(trie_keys, next_word_ids):
# def _clean_trie_key(trie, trie_key):
# def _trie_scores_py_func(trie, beam_search=False):
#     def _py_func(log_probs, sequence_lengths, scores, trie_keys, trie_exclude):
#     def _py_func(logits, trie_keys, trie_exclude):
#       def _py_func_inner(batch_idx):
# def _get_valid_continuation_idxs(trie, trie_exclude_batch_set, trie_key):
# def _create_trie_exclude_batch_set(trie, trie_exclude_batch):
# def _get_scores(log_probs, sequence_lengths, length_penalty_weight):
# def _length_penalty(sequence_lengths, penalty_factor):
# def _get_trie_scores(log_probs, sequence_lengths, length_penalty_weight,
#                      trie_keys, trie_exclude, trie):
# def _mask_probs(probs, eos_token, finished):
# def _check_maybe(t):
# def _maybe_tensor_gather_helper(gather_indices, gather_from, batch_size,
#                                 range_size, gather_shape):
# def _tensor_gather_helper(gather_indices,
#                           gather_from,
#                           batch_size,
#                           range_size,
#                           gather_shape,
#                           name=None):


# Path: px/nmt/utils/trie_decoder_utils_test.py
import os.path
import shutil
import tempfile
import tensorflow as tf
from px.nmt.utils import trie_decoder_utils
from px.proto import aqa_pb2

        ('a b c b a c', 'y'),
        ('a c c b', 'z'),
    ]

    with open(cls.vocab_path, 'w') as f:
      f.write('\n'.join(cls.test_vocab))


  @classmethod
  def tearDownClass(cls):
    shutil.rmtree(cls.tmpdir)


  def testSaveLoadTrie(self):
    trie = trie_decoder_utils.DecoderTrie(TrieDecoderUtilsTest.vocab_path)
    for question, answer in TrieDecoderUtilsTest.test_trie_entries:
      trie.insert_question(question, answer)
    NEW_KEY = 'newkey'
    self.assertNotIn(NEW_KEY, trie)
    self.assertFalse(os.path.exists(TrieDecoderUtilsTest.save_path))
    trie.save_to_file(TrieDecoderUtilsTest.save_path)
    self.assertTrue(os.path.exists(TrieDecoderUtilsTest.save_path))
    trie[NEW_KEY] = 'foo'
    self.assertIn(NEW_KEY, trie)
    loaded_trie = trie_decoder_utils.DecoderTrie.load_from_file(
        TrieDecoderUtilsTest.save_path)
    self.assertNotIn(NEW_KEY, loaded_trie)

  def testCreateEmptyTrie(self):
    eos_token = '</s>'

    trie = trie_decoder_utils.DecoderTrie(TrieDecoderUtilsTest.vocab_path,