def __init__(self): super().__init__() self.vocab = Vocabulary(input_vocab, unk_idx=input_vocab.idx[UNK]) self.model = traced_model self.output_layer = output_layer self.pad_idx = jit.Attribute(input_vocab.idx[PAD], int)
def __init__( self, embedding, jit_module, word_vocab, dict_vocab, action_vocab, word_unk_idx=0, dict_unk_idx=0, ): super().__init__() self.word_vocab = Vocabulary(word_vocab.itos, unk_idx=word_unk_idx) self.dict_vocab = Vocabulary(dict_vocab.itos, unk_idx=dict_unk_idx) self.action_vocab = Vocabulary(action_vocab.itos, unk_idx=-1) self.embedding = embedding self.jit_module = jit_module
def __init__(self): super().__init__() self.vocab = Vocabulary(input_vocab, unk_idx=input_vocab.idx[UNK]) self.max_byte_len = jit.Attribute(max_byte_len, int) self.byte_offset_for_non_padding = jit.Attribute( byte_offset_for_non_padding, int) self.pad_idx = jit.Attribute(input_vocab.idx[PAD], int) self.model = traced_model self.output_layer = output_layer
def setUp(self): vocab_list = ["UNK", "a", "b", "c", "d"] self.vocab = Vocabulary(vocab_list)
def test_custom_unk(self): vocab_list = ["a", "UNK", "b", "c", "d"] vocab = Vocabulary(vocab_list, unk_idx=1) self.assertEqual([0, 1, 3, 4], vocab.lookup_indices_1d(["a", "e", "c", "d"]))