def __init__(self, type_context: TypeContext):
super().__init__()
self.lookup = torch.zeros(type_context.get_object_count(), 3)
self.lookup[0] = torch.Tensor([1, 1, 1])
self.lookup[1] = torch.Tensor([-1, -1, -1])
self.lookup[2] = torch.Tensor([3, 3, 3])
self.lookup[3] = torch.Tensor([10, 0, 1])
self.lookup[6] = torch.Tensor([3, 5, 1])
def __init__(
self,
rnn_cell: TreeRNNCell,
action_selector: ActionSelector,
type_vectorizer: VectorizerBase,
type_context: TypeContext #,
#bce_pos_weight=1.0
):
super().__init__()
self.rnn_cell = rnn_cell
self.action_selector = action_selector
self.type_vectorizer = type_vectorizer
self.type_context = type_context
self.object_embeddings = nn.Embedding(type_context.get_object_count(),
rnn_cell.hidden_size)
def get_default_nonretrieval_decoder(type_context: TypeContext,
rnn_hidden_size: int) -> TreeDecoder:
object_vectorizer = vectorizers.TorchDeepEmbed(
type_context.get_object_count(), rnn_hidden_size)
ast_embed_size = int(rnn_hidden_size / 2)
type_vectorizer = vectorizers.TorchDeepEmbed(type_context.get_type_count(),
ast_embed_size)
rnn_cell = TreeRNNCellLSTM(ast_embed_size, rnn_hidden_size)
#rnn_cell = TreeCellOnlyAttn(rnn_hidden_size, rnn_hidden_size)
#rnn_cell = TreeRNNCellGRU(rnn_hidden_size, rnn_hidden_size)
action_selector = SimpleActionSelector(
rnn_cell.output_size,
objectselector.get_default_object_selector(type_context,
object_vectorizer),
type_context)
return TreeRNNDecoder(rnn_cell, action_selector, type_vectorizer,
type_context)
