class Graph2HybridMetric(CodeGraphModel): @staticmethod def _init_metadata(hyperparameters: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Graph2HybridMetric, Graph2HybridMetric)._init_metadata(hyperparameters, raw_metadata) TypeClassificationModel._init_metadata(hyperparameters, raw_metadata) @staticmethod def _load_metadata_from_sample(hyperparameters: Dict[str, Any], raw_sample: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Graph2HybridMetric, Graph2HybridMetric)._load_metadata_from_sample(hyperparameters, raw_sample, raw_metadata) TypeClassificationModel._load_metadata_from_sample(hyperparameters, raw_sample, raw_metadata) def _finalise_metadata(self, raw_metadata_list: List[Dict[str, Any]]) -> Dict[str, Any]: final_metadata = super()._finalise_metadata(raw_metadata_list) self.__type_classification._finalise_metadata(raw_metadata_list, final_metadata) return final_metadata @staticmethod def get_default_hyperparameters() -> Dict[str, Any]: defaults = CodeGraphModel.get_default_hyperparameters() defaults.update({ 'margin': 2, 'max_type_annotation_vocab_size': 100, 'strip_type_parameters': True }) return defaults def __init__(self, hyperparameters, run_name: Optional[str] = None, model_save_dir: Optional[str] = None, log_save_dir: Optional[str] = None): super().__init__(hyperparameters, run_name, model_save_dir, log_save_dir) self.__type_metric = TypeMetricLearningModel(self, type_representation_size=self.hyperparameters['cg_ggnn_hidden_size'], margin=self.hyperparameters['margin']) self.__type_classification = TypeClassificationModel(self) def _make_placeholders(self, is_train: bool) -> None: super()._make_placeholders(is_train) self.placeholders['typed_annotation_node_ids'] = tf.placeholder(tf.int32, shape=( None,), name="typed_annotation_node_ids") self.__type_metric._make_placeholders(is_train) self.__type_classification._make_placeholders(is_train) def _make_parameters(self): super()._make_parameters() representation_size = self.hyperparameters['cg_ggnn_hidden_size'] self.parameters['metric_to_classification_linear_layer'] = \ tf.get_variable(name='metric_to_classification_linear_layer', initializer=tf.random_normal_initializer(), shape=[representation_size, representation_size], ) self.__type_classification._make_parameters(representation_size=representation_size) def _make_model(self, is_train: bool = True): super()._make_model(is_train) # Gather up CG node representations for the nodes we want to type target_node_representations_from_cg = tf.gather(params=self.ops['cg_node_representations'], indices=self.placeholders['typed_annotation_node_ids']) self.ops['target_representations'] = target_node_representations_from_cg self.__type_metric._make_model(target_node_representations_from_cg, is_train) classification_representation = tf.matmul(target_node_representations_from_cg, self.parameters['metric_to_classification_linear_layer']) classification_representation = tf.nn.dropout(classification_representation, rate=1-self.placeholders['dropout_keep_rate']) self.__type_classification._make_model(classification_representation, is_train) @staticmethod def _load_data_from_sample(hyperparameters: Dict[str, Any], metadata: Dict[str, Any], raw_sample: Dict[str, Any], result_holder: Dict[str, Any], is_train: bool = True) -> bool: keep_sample = super(Graph2HybridMetric, Graph2HybridMetric)._load_data_from_sample( hyperparameters, metadata, raw_sample, result_holder, is_train) if not keep_sample: return False target_node_idxs, target_class, target_class_id = [], [], [] for node_idx, annotation_data in raw_sample['supernodes'].items(): node_idx = int(node_idx) annotation = annotation_data['annotation'] if is_train and ignore_type_annotation(annotation): continue target_node_idxs.append(node_idx) target_class.append(annotation) target_class_id.append(TypeClassificationModel._get_idx_for_type(annotation, metadata, hyperparameters)) result_holder['target_node_idxs'] = np.array(target_node_idxs, dtype=np.uint16) result_holder['target_type'] = target_class result_holder['variable_target_class'] = np.array(target_class_id, dtype=np.uint16) return len(target_node_idxs) > 0 def _init_minibatch(self, batch_data: Dict[str, Any]) -> None: super()._init_minibatch(batch_data) batch_data['batch_target_node_idxs'] = [] self.__type_metric._init_minibatch(batch_data) self.__type_classification._init_minibatch(batch_data) def _extend_minibatch_by_sample(self, batch_data: Dict[str, Any], sample: Dict[str, Any]) -> bool: batch_finished = super()._extend_minibatch_by_sample(batch_data, sample) # The offset has already been updated in our parent... original_cg_node_offset = batch_data['cg_node_offset'] - sample['num_nodes'] batch_data['batch_target_node_idxs'].extend(sample['target_node_idxs'] + original_cg_node_offset) self.__type_classification._extend_minibatch_by_sample(batch_data, sample) self.__type_metric._extend_minibatch_by_sample(batch_data, sample) return batch_finished def _finalise_minibatch(self, batch_data: Dict[str, Any], is_train: bool) -> Dict[tf.Tensor, Any]: minibatch = super()._finalise_minibatch(batch_data, is_train) write_to_minibatch( minibatch, self.placeholders['typed_annotation_node_ids'], batch_data['batch_target_node_idxs']) self.__type_classification._finalise_minibatch(batch_data, is_train, minibatch) self.__type_metric._finalise_minibatch(batch_data, is_train, minibatch) return minibatch def create_index(self, data_paths: List[RichPath]): self.__type_metric.create_index(data_paths, self.metadata) # ------- These are the bits that we only need for test-time: def _encode_one_test_sample(self, sample_data_dict: Dict[tf.Tensor, Any]) -> Tuple[tf.Tensor, Optional[tf.Tensor]]: return (self.sess.run(self.ops['target_representations'], feed_dict=sample_data_dict), None) def annotate_single(self, raw_sample: Dict[str, Any], loaded_test_sample: Dict[str, Any], provenance: str): s_t = time.time() a = self.__type_metric.annotate_single(raw_sample, loaded_test_sample, provenance, self.metadata) print("single prediction in %f sec." % (time.time() - s_t)) return a
class Sequence2HybridMetric(Model): @staticmethod def get_default_hyperparameters() -> Dict[str, Any]: defaults = Model.get_default_hyperparameters() defaults.update({ 'token_embedding_style': 'Subtoken', # One of ['Token', 'CharCNN', 'Subtoken'] 'token_vocab_size': 10000, # Used only when cg_node_label_embedding_style is 'Token' or 'Subtoken' 'token_char_length': 16, # Used only when cg_node_label_embedding_style is 'CharCNN' 'token_max_subtokens': 5, # Used only when cg_node_label_embedding_style is 'Subtoken' 'token_embedding_size': 64, 'seq_layer_type': 'birnn', # One of ['BiRNN', 'SelfAtt'] 'num_seq_layers': 2, 'use_consistency_layer': True, 'max_seq_len': 50000, # BiRNN-specific 'birnn_hidden_size': 32, #Transformer-specific 'tranformer_hidden_size': 64, 'tranformer_num_attention_heads': 8, 'tranformer_intermediate_size': 128, 'batch_size': 32, 'margin': 2, 'max_type_annotation_vocab_size': 100, 'strip_type_parameters': True }) return defaults def __init__(self, hyperparameters, run_name: Optional[str] = None, model_save_dir: Optional[str] = None, log_save_dir: Optional[str] = None): super().__init__(hyperparameters, run_name, model_save_dir, log_save_dir) seq_layer_type = self.hyperparameters['seq_layer_type'].lower() if seq_layer_type == 'birnn': out_size = 2 * self.hyperparameters['birnn_hidden_size'] elif seq_layer_type == 'selfatt': out_size = self.hyperparameters['tranformer_hidden_size'] else: raise ValueError('Unrecognized type of Sequential Layer %s' % seq_layer_type) self.__type_metric = TypeMetricLearningModel( self, type_representation_size=out_size, margin=self.hyperparameters['margin']) self.__type_classification = TypeClassificationModel(self) def _make_parameters(self): super()._make_parameters() TokenEmbedder.make_parameters('token', self.parameters, self.metadata, self.hyperparameters) seq_layer_type = self.hyperparameters['seq_layer_type'].lower() if seq_layer_type == 'birnn': self.parameters['seq_layers'] = [ tf.keras.layers.Bidirectional( tf.keras.layers.GRU( self.hyperparameters['birnn_hidden_size'], return_sequences=True)) for _ in range(self.hyperparameters['num_seq_layers']) ] out_size = 2 * self.hyperparameters['birnn_hidden_size'] elif seq_layer_type == 'selfatt': self.parameters['seq_layers'] = [ lambda input: transformer_model( input, hidden_size=self.hyperparameters['tranformer_hidden_size'], num_hidden_layers=1, num_attention_heads=self.hyperparameters[ 'tranformer_num_attention_heads'], intermediate_size=self.hyperparameters[ 'tranformer_intermediate_size']) for _ in range(self.hyperparameters['num_seq_layers']) ] out_size = self.hyperparameters['tranformer_hidden_size'] else: raise ValueError('Unrecognized type of Sequential Layer %s' % seq_layer_type) self.parameters['metric_to_classification_linear_layer'] = \ tf.get_variable(name='metric_to_classification_linear_layer', initializer=tf.random_normal_initializer(), shape=[out_size, out_size], ) self.__type_classification._make_parameters( representation_size=out_size) def _make_placeholders(self, is_train: bool) -> None: super()._make_placeholders(is_train) TokenEmbedder.make_placeholders('token', self.placeholders, hyperparameters=self.hyperparameters) self.placeholders['embedding_gather_matrix'] = tf.placeholder( tf.int32, shape=(None, None), name='embedding_gather_matrix') self.placeholders['sequence_lengths'] = tf.placeholder( tf.int32, shape=(None, ), name="sequence_lengths") self.placeholders['variable_bound_token_ids'] = tf.placeholder( tf.int32, shape=(None, ), name="variable_bound_token_ids") self.placeholders['token_variable_ids'] = tf.placeholder( tf.int32, shape=(None, ), name="token_variable_ids") self.placeholders['num_variables'] = tf.placeholder( tf.int32, shape=[], name="num_variables") self.__type_metric._make_placeholders(is_train) self.__type_classification._make_placeholders(is_train) def _make_model(self, is_train: bool = True): super()._make_model(is_train) initial_token_embeddings = TokenEmbedder.make_model( 'token', self.placeholders, self.parameters, self.hyperparameters, is_train) # T x D sequence_token_embeddings = tf.gather_nd( params=initial_token_embeddings, indices=tf.expand_dims( self.placeholders['embedding_gather_matrix'], axis=-1)) # B x max-len x D def get_variable_embeddings(all_sequence_embeddings): flat_sequence_embeddings = tf.reshape( all_sequence_embeddings, (-1, all_sequence_embeddings.get_shape()[-1])) # B*max-len x D target_token_embeddings = tf.gather( params=flat_sequence_embeddings, indices=self.placeholders['variable_bound_token_ids']) return tf.unsorted_segment_mean( data=target_token_embeddings, segment_ids=self.placeholders['token_variable_ids'], num_segments=self.placeholders[ 'num_variables'] # TODO: Do not depend in any way on the classes. ) # num-variables x H # Multiple layers of BiRNN/Transformer and "consistency" layer. current_out = sequence_token_embeddings # B x max-len x H for i, seq_layer in enumerate(self.parameters['seq_layers']): # Mask out-of-sequence-tokens mask = tf.cast( tf.reshape(tf.range(tf.shape(sequence_token_embeddings)[1]), [1, -1, 1]) < tf.reshape(self.placeholders['sequence_lengths'], [-1, 1, 1]), current_out.dtype) current_out *= mask with tf.variable_scope('seqlayer_%s' % i): current_out = seq_layer(current_out) if i < len(self.parameters['seq_layers'] ) - 1 and self.hyperparameters['use_consistency_layer']: # Apply "consistency" layer to all layers but the last variable_embeddings = get_variable_embeddings( current_out) # num-variables x H variable_embedding_per_token = tf.gather( params=variable_embeddings, indices=self.placeholders['token_variable_ids'] ) # num-variable-tokens x H current_out_shape = tf.shape(current_out) updates = tf.scatter_nd(indices=tf.expand_dims( self.placeholders['variable_bound_token_ids'], axis=-1), updates=variable_embedding_per_token, shape=[ current_out_shape[0] * current_out_shape[1], tf.shape(current_out)[2] ]) current_out += tf.reshape(updates, tf.shape(current_out)) variable_embeddings = get_variable_embeddings(current_out) self.ops['variable_embeddings'] = variable_embeddings self.__type_metric._make_model(variable_embeddings, is_train) classification_representation = tf.matmul( variable_embeddings, self.parameters['metric_to_classification_linear_layer']) classification_representation = tf.nn.dropout( classification_representation, rate=1 - self.placeholders['dropout_keep_rate']) self.__type_classification._make_model(classification_representation, is_train) @staticmethod def _init_metadata(hyperparameters: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Sequence2HybridMetric, Sequence2HybridMetric)._init_metadata(hyperparameters, raw_metadata) TokenEmbedder.init_metadata('token', raw_metadata, hyperparameters) TypeClassificationModel._init_metadata(hyperparameters, raw_metadata) @staticmethod def _load_metadata_from_sample(hyperparameters: Dict[str, Any], raw_sample: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Sequence2HybridMetric, Sequence2HybridMetric)._load_metadata_from_sample( hyperparameters, raw_sample, raw_metadata) TokenEmbedder.load_metadata_from_sample('token', [ raw_sample['nodes'][i] for i in raw_sample['token-sequence'] if i < len(raw_sample['nodes']) ], raw_metadata, hyperparameters) TypeClassificationModel._load_metadata_from_sample( hyperparameters, raw_sample, raw_metadata) def _finalise_metadata( self, raw_metadata_list: List[Dict[str, Any]]) -> Dict[str, Any]: final_metadata = super()._finalise_metadata(raw_metadata_list) TokenEmbedder.finalise_metadata('token', raw_metadata_list, final_metadata, self.hyperparameters) self.__type_classification._finalise_metadata(raw_metadata_list, final_metadata) return final_metadata @staticmethod def _load_data_from_sample(hyperparameters: Dict[str, Any], metadata: Dict[str, Any], raw_sample: Dict[str, Any], result_holder: Dict[str, Any], is_train: bool = True) -> bool: keep_sample = super(Sequence2HybridMetric, Sequence2HybridMetric)._load_data_from_sample( hyperparameters, metadata, raw_sample, result_holder, is_train) if not keep_sample: return False token_node_idxs = set(raw_sample['token-sequence']) node_idx_to_supernode_idx = {} # type: Dict[int, int] for from_idx, to_idxs in raw_sample['edges']['OCCURRENCE_OF'].items(): from_idx = int(from_idx) if from_idx not in token_node_idxs: # Some supernodes do not have an associated token. Such nodes are attributes if str(from_idx) in raw_sample['edges']['CHILD']: right_token_idx = max( raw_sample['edges']['CHILD'][str(from_idx)]) assert right_token_idx in token_node_idxs from_idx = right_token_idx else: continue for to_idx in to_idxs: node_idx_to_supernode_idx[from_idx] = to_idx supernodes_with_related_nodes = set(node_idx_to_supernode_idx.values()) variable_types = [] # type: List[str] variable_type_idxs = [] # type: List[int] ignored_supernodes = set() supernode_idxs_to_annotated_variable_idx = {} # type: Dict[int, int] for node_idx, supernode_data in raw_sample['supernodes'].items(): node_idx = int(node_idx) annotation = supernode_data['annotation'] if ignore_type_annotation(annotation) and is_train: ignored_supernodes.add(node_idx) continue if node_idx not in supernodes_with_related_nodes: ignored_supernodes.add(node_idx) continue variable_idx = len(supernode_idxs_to_annotated_variable_idx) variable_types.append(annotation) variable_type_idxs.append( TypeClassificationModel._get_idx_for_type( annotation, metadata, hyperparameters)) supernode_idxs_to_annotated_variable_idx[node_idx] = variable_idx if len(variable_types) == 0: return False token_idx, variable_idx = [], [] def create_token_sequence(): for i, node_idx in enumerate(raw_sample['token-sequence']): supernode_idx = node_idx_to_supernode_idx.get(node_idx) if supernode_idx is not None: annotated_variable_idxs = supernode_idxs_to_annotated_variable_idx.get( supernode_idx) if annotated_variable_idxs is not None: token_idx.append(i) variable_idx.append(annotated_variable_idxs) yield raw_sample['nodes'][node_idx] token_sequence = list(create_token_sequence()) if len(token_sequence) > hyperparameters['max_seq_len']: return False # Did we see at least one token per variable? assert len(np.unique(variable_idx)) == len(variable_types) TokenEmbedder.load_data_from_sample('token', metadata, token_sequence, result_holder, hyperparameters, is_train) result_holder['sequence_length'] = len(token_sequence) result_holder['variable_token_idxs'] = np.array(token_idx, dtype=np.uint32) result_holder['variable_idxs'] = np.array(variable_idx, dtype=np.uint32) result_holder['target_type'] = variable_types result_holder['variable_target_class'] = np.array(variable_type_idxs, dtype=np.uint32) result_holder['ignored_supernodes'] = ignored_supernodes return keep_sample def _init_minibatch(self, batch_data: Dict[str, Any]) -> None: super()._init_minibatch(batch_data) TokenEmbedder.init_minibatch('token', batch_data, self.hyperparameters) batch_data['batch_sequence_lengths'] = [] batch_data['batch_variable_token_idxs'] = [] batch_data['batch_variable_idxs'] = [] self.__type_metric._init_minibatch(batch_data) self.__type_classification._init_minibatch(batch_data) def _extend_minibatch_by_sample(self, batch_data: Dict[str, Any], sample: Dict[str, Any]) -> bool: super()._extend_minibatch_by_sample(batch_data, sample) TokenEmbedder.extend_minibatch_by_sample('token', batch_data, sample, self.hyperparameters) batch_data['batch_sequence_lengths'].append(sample['sequence_length']) batch_data['batch_variable_token_idxs'].append( sample['variable_token_idxs']) batch_data['batch_variable_idxs'].extend( sample['variable_idxs'] + len(batch_data['batch_target_variable_type'])) self.__type_metric._extend_minibatch_by_sample(batch_data, sample) self.__type_classification._extend_minibatch_by_sample( batch_data, sample) return len(batch_data['batch_sequence_lengths'] ) >= self.hyperparameters['batch_size'] def _finalise_minibatch(self, batch_data: Dict[str, Any], is_train: bool) -> Dict[tf.Tensor, Any]: minibatch = super()._finalise_minibatch(batch_data, is_train) TokenEmbedder.finalise_minibatch('token', batch_data, self.placeholders, minibatch, self.hyperparameters, is_train) max_sequence_len = max(batch_data['batch_sequence_lengths']) batch_size = len(batch_data['batch_sequence_lengths']) # TODO: Later, force split very long sequences # embedding_gather_matrix contains the idxs such that embeddings[embedding_gather_matrix] creates a B x max_sequence_len x D matrix embedding_gather_matrix = np.zeros((batch_size, max_sequence_len), dtype=np.int32) current_idx = 0 for i, length in enumerate(batch_data['batch_sequence_lengths']): embedding_gather_matrix[i, :length] = np.arange(start=current_idx, stop=current_idx + length) current_idx += length write_to_minibatch(minibatch, self.placeholders['embedding_gather_matrix'], embedding_gather_matrix) total_variable_tokens = len(batch_data['batch_variable_idxs']) variable_token_gather_idxs = np.empty(total_variable_tokens, dtype=np.int32) gathered_tokens_so_far = 0 for i, tidx in enumerate(batch_data['batch_variable_token_idxs']): num_variable_tokens = len(tidx) variable_token_gather_idxs[ gathered_tokens_so_far:gathered_tokens_so_far + num_variable_tokens] = tidx + i * max_sequence_len gathered_tokens_so_far += num_variable_tokens write_to_minibatch(minibatch, self.placeholders['variable_bound_token_ids'], variable_token_gather_idxs) write_to_minibatch(minibatch, self.placeholders['sequence_lengths'], batch_data['batch_sequence_lengths']) write_to_minibatch(minibatch, self.placeholders['token_variable_ids'], batch_data['batch_variable_idxs']) self.__type_metric._finalise_minibatch(batch_data, is_train, minibatch) self.__type_classification._finalise_minibatch(batch_data, is_train, minibatch) minibatch[self.placeholders['num_variables']] = len( batch_data['batch_target_variable_type']) return minibatch def create_index(self, data_paths: List[RichPath]): self.__type_metric.create_index(data_paths, self.metadata) # ------- These are the bits that we only need for test-time: def _encode_one_test_sample( self, sample_data_dict: Dict[tf.Tensor, Any] ) -> Tuple[tf.Tensor, Optional[tf.Tensor]]: return (self.sess.run(self.ops['target_representations'], feed_dict=sample_data_dict), None) def annotate_single(self, raw_sample: Dict[str, Any], loaded_test_sample: Dict[str, Any], provenance: str): return self.__type_metric.annotate_single(raw_sample, loaded_test_sample, provenance, self.metadata)
class Path2Annotation(PathBasedModel): @staticmethod def get_default_hyperparameters() -> Dict[str, Any]: defaults = PathBasedModel.get_default_hyperparameters() defaults.update({ 'max_type_annotation_vocab_size': 100, }) return defaults def __init__(self, hyperparameters, run_name: Optional[str] = None, model_save_dir: Optional[str] = None, log_save_dir: Optional[str] = None): super().__init__(hyperparameters, run_name, model_save_dir, log_save_dir) self.__type_classification = TypeClassificationModel(self) def _make_parameters(self): super()._make_parameters() self.__type_classification._make_parameters( representation_size=self.hyperparameters['path_encoding_size']) def _make_placeholders(self, is_train: bool) -> None: super()._make_placeholders(is_train) self.placeholders['typed_annotation_node_ids'] = tf.placeholder( tf.int32, shape=(None, ), name="typed_annotation_node_ids") self.__type_classification._make_placeholders(is_train) def _make_model(self, is_train: bool = True): super()._make_model(is_train) self.__type_classification._make_model( self.ops['target_variable_representations'], is_train) @staticmethod def _init_metadata(hyperparameters: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Path2Annotation, Path2Annotation)._init_metadata(hyperparameters, raw_metadata) TypeClassificationModel._init_metadata(hyperparameters, raw_metadata) @staticmethod def _load_metadata_from_sample(hyperparameters: Dict[str, Any], raw_sample: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Path2Annotation, Path2Annotation)._load_metadata_from_sample( hyperparameters, raw_sample, raw_metadata) TypeClassificationModel._load_metadata_from_sample( hyperparameters, raw_sample, raw_metadata) def _finalise_metadata( self, raw_metadata_list: List[Dict[str, Any]]) -> Dict[str, Any]: final_metadata = super()._finalise_metadata(raw_metadata_list) self.__type_classification._finalise_metadata(raw_metadata_list, final_metadata) return final_metadata @staticmethod def _load_data_from_sample(hyperparameters: Dict[str, Any], metadata: Dict[str, Any], raw_sample: Dict[str, Any], result_holder: Dict[str, Any], is_train: bool = True) -> bool: keep_sample = super(Path2Annotation, Path2Annotation)._load_data_from_sample( hyperparameters, metadata, raw_sample, result_holder, is_train) if not keep_sample: return False target_class = [] for node_idx, annotation_data in raw_sample['supernodes'].items(): annotation = annotation_data['annotation'] if is_train and ignore_type_annotation(annotation): continue target_class.append( TypeClassificationModel._get_idx_for_type( annotation, metadata, hyperparameters)) result_holder['variable_target_class'] = np.array(target_class, dtype=np.uint16) return len(target_class) > 0 def _init_minibatch(self, batch_data: Dict[str, Any]) -> None: super()._init_minibatch(batch_data) self.__type_classification._init_minibatch(batch_data) def _extend_minibatch_by_sample(self, batch_data: Dict[str, Any], sample: Dict[str, Any]) -> bool: batch_finished = super()._extend_minibatch_by_sample( batch_data, sample) self.__type_classification._extend_minibatch_by_sample( batch_data, sample) return batch_finished def _finalise_minibatch(self, batch_data: Dict[str, Any], is_train: bool) -> Dict[tf.Tensor, Any]: minibatch = super()._finalise_minibatch(batch_data, is_train) self.__type_classification._finalise_minibatch(batch_data, is_train, minibatch) return minibatch # ------- These are the bits that we only need for test-time: def _encode_one_test_sample( self, sample_data_dict: Dict[tf.Tensor, Any] ) -> Tuple[tf.Tensor, Optional[tf.Tensor]]: return (self.sess.run(self.ops['target_representations'], feed_dict=sample_data_dict), None) def class_id_to_class(self, class_id: int) -> str: return self.__type_classification.class_id_to_class(class_id) def annotate_single(self, raw_sample: Dict[str, Any], loaded_test_sample: Dict[str, Any], provenance: str): return self.__type_classification.annotate_single( raw_sample, loaded_test_sample, provenance)
class Graph2Annotation(CodeGraphModel): @staticmethod def get_default_hyperparameters() -> Dict[str, Any]: defaults = CodeGraphModel.get_default_hyperparameters() defaults.update({ 'max_type_annotation_vocab_size': 100, }) return defaults def __init__(self, hyperparameters, run_name: Optional[str] = None, model_save_dir: Optional[str] = None, log_save_dir: Optional[str] = None): super().__init__(hyperparameters, run_name, model_save_dir, log_save_dir) self.__type_classification = TypeClassificationModel(self) def _make_parameters(self): super()._make_parameters() self.__type_classification._make_parameters(representation_size=self.hyperparameters['cg_ggnn_hidden_size']) def _make_placeholders(self, is_train: bool) -> None: super()._make_placeholders(is_train) self.placeholders['typed_annotation_node_ids'] = tf.placeholder(tf.int32, shape=( None,), name="typed_annotation_node_ids") self.__type_classification._make_placeholders(is_train) def _make_model(self, is_train: bool = True): super()._make_model(is_train) # Gather up CG node representations for the nodes we want to type target_node_representations_from_cg = tf.gather(params=self.ops['cg_node_representations'], indices=self.placeholders['typed_annotation_node_ids']) self.ops['target_representations'] = target_node_representations_from_cg self.__type_classification._make_model( target_node_representations_from_cg, is_train) @staticmethod def _init_metadata(hyperparameters: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Graph2Annotation, Graph2Annotation)._init_metadata( hyperparameters, raw_metadata) TypeClassificationModel._init_metadata(hyperparameters, raw_metadata) @staticmethod def _load_metadata_from_sample(hyperparameters: Dict[str, Any], raw_sample: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Graph2Annotation, Graph2Annotation)._load_metadata_from_sample( hyperparameters, raw_sample, raw_metadata) TypeClassificationModel._load_metadata_from_sample(hyperparameters, raw_sample, raw_metadata) def _finalise_metadata(self, raw_metadata_list: List[Dict[str, Any]]) -> Dict[str, Any]: final_metadata = super()._finalise_metadata(raw_metadata_list) self.__type_classification._finalise_metadata(raw_metadata_list, final_metadata) return final_metadata @staticmethod def _load_data_from_sample(hyperparameters: Dict[str, Any], metadata: Dict[str, Any], raw_sample: Dict[str, Any], result_holder: Dict[str, Any], is_train: bool = True) -> bool: keep_sample = super(Graph2Annotation, Graph2Annotation)._load_data_from_sample( hyperparameters, metadata, raw_sample, result_holder, is_train) if not keep_sample: return False target_node_idxs, target_class = [], [] for node_idx, annotation_data in raw_sample['supernodes'].items(): node_idx = int(node_idx) annotation = annotation_data['annotation'] if is_train and ignore_type_annotation(annotation): continue target_node_idxs.append(node_idx) target_class.append(TypeClassificationModel._get_idx_for_type(annotation, metadata, hyperparameters)) result_holder['target_node_idxs'] = np.array(target_node_idxs, dtype=np.uint16) result_holder['variable_target_class'] = np.array(target_class, dtype=np.uint16) return len(target_node_idxs) > 0 def _init_minibatch(self, batch_data: Dict[str, Any]) -> None: super()._init_minibatch(batch_data) batch_data['batch_target_node_idxs'] = [] self.__type_classification._init_minibatch(batch_data) def _extend_minibatch_by_sample(self, batch_data: Dict[str, Any], sample: Dict[str, Any]) -> bool: batch_finished = super()._extend_minibatch_by_sample(batch_data, sample) # The offset has already been updated in our parent... original_cg_node_offset = batch_data['cg_node_offset'] - sample['num_nodes'] batch_data['batch_target_node_idxs'].extend(sample['target_node_idxs'] + original_cg_node_offset) self.__type_classification._extend_minibatch_by_sample(batch_data, sample) return batch_finished def _finalise_minibatch(self, batch_data: Dict[str, Any], is_train: bool) -> Dict[tf.Tensor, Any]: minibatch = super()._finalise_minibatch(batch_data, is_train) write_to_minibatch( minibatch, self.placeholders['typed_annotation_node_ids'], batch_data['batch_target_node_idxs']) self.__type_classification._finalise_minibatch(batch_data, is_train, minibatch) return minibatch # ------- These are the bits that we only need for test-time: def _encode_one_test_sample(self, sample_data_dict: Dict[tf.Tensor, Any]) -> Tuple[tf.Tensor, Optional[tf.Tensor]]: return (self.sess.run(self.ops['target_representations'], feed_dict=sample_data_dict), None) def class_id_to_class(self, class_id: int) -> str: return self.__type_classification.class_id_to_class(class_id) def annotate_single(self, raw_sample: Dict[str, Any], loaded_test_sample: Dict[str, Any], provenance: str): return self.__type_classification.annotate_single(raw_sample, loaded_test_sample, provenance)
class Path2HybridMetric(PathBasedModel): @staticmethod def _init_metadata(hyperparameters: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Path2HybridMetric, Path2HybridMetric)._init_metadata(hyperparameters, raw_metadata) TypeClassificationModel._init_metadata(hyperparameters, raw_metadata) @staticmethod def _load_metadata_from_sample(hyperparameters: Dict[str, Any], raw_sample: Dict[str, Any], raw_metadata: Dict[str, Any]) -> None: super(Path2HybridMetric, Path2HybridMetric)._load_metadata_from_sample( hyperparameters, raw_sample, raw_metadata) TypeClassificationModel._load_metadata_from_sample( hyperparameters, raw_sample, raw_metadata) def _finalise_metadata( self, raw_metadata_list: List[Dict[str, Any]]) -> Dict[str, Any]: final_metadata = super()._finalise_metadata(raw_metadata_list) self.__type_classification._finalise_metadata(raw_metadata_list, final_metadata) return final_metadata @staticmethod def get_default_hyperparameters() -> Dict[str, Any]: defaults = PathBasedModel.get_default_hyperparameters() defaults.update({ 'margin': 2, 'max_type_annotation_vocab_size': 100, 'strip_type_parameters': True }) return defaults def __init__(self, hyperparameters, run_name: Optional[str] = None, model_save_dir: Optional[str] = None, log_save_dir: Optional[str] = None): super().__init__(hyperparameters, run_name, model_save_dir, log_save_dir) self.__type_classification = TypeClassificationModel(self) self.__type_metric = TypeMetricLearningModel( self, type_representation_size=self. hyperparameters['path_encoder_lstm_hidden_size'], margin=self.hyperparameters['margin']) self.__type_classification = TypeClassificationModel(self) def _make_parameters(self): super()._make_parameters() representation_size = self.hyperparameters['path_encoding_size'] self.parameters['metric_to_classification_linear_layer'] = \ tf.get_variable(name='metric_to_classification_linear_layer', initializer=tf.random_normal_initializer(), shape=[representation_size, representation_size], ) self.__type_classification._make_parameters( representation_size=representation_size) def _make_placeholders(self, is_train: bool) -> None: super()._make_placeholders(is_train) self.__type_metric._make_placeholders(is_train) self.__type_classification._make_placeholders(is_train) def _make_model(self, is_train: bool = True): super()._make_model(is_train) self.__type_metric._make_model( self.ops['target_variable_representations'], is_train) classification_representation = tf.matmul( self.ops['target_variable_representations'], self.parameters['metric_to_classification_linear_layer']) classification_representation = tf.nn.dropout( classification_representation, rate=1 - self.placeholders['dropout_keep_rate']) self.__type_classification._make_model(classification_representation, is_train) @staticmethod def _load_data_from_sample(hyperparameters: Dict[str, Any], metadata: Dict[str, Any], raw_sample: Dict[str, Any], result_holder: Dict[str, Any], is_train: bool = True) -> bool: keep_sample = super(Path2HybridMetric, Path2HybridMetric)._load_data_from_sample( hyperparameters, metadata, raw_sample, result_holder, is_train) if not keep_sample: return False target_class, target_class_id = [], [] for node_idx, annotation_data in raw_sample['supernodes'].items(): annotation = annotation_data['annotation'] if is_train and ignore_type_annotation(annotation): continue target_class.append(annotation) target_class_id.append( TypeClassificationModel._get_idx_for_type( annotation, metadata, hyperparameters)) result_holder['target_type'] = target_class result_holder['variable_target_class'] = np.array(target_class_id, dtype=np.uint16) return len(target_class) > 0 def _init_minibatch(self, batch_data: Dict[str, Any]) -> None: super()._init_minibatch(batch_data) self.__type_metric._init_minibatch(batch_data) self.__type_classification._init_minibatch(batch_data) def _extend_minibatch_by_sample(self, batch_data: Dict[str, Any], sample: Dict[str, Any]) -> bool: batch_finished = super()._extend_minibatch_by_sample( batch_data, sample) self.__type_metric._extend_minibatch_by_sample(batch_data, sample) self.__type_classification._extend_minibatch_by_sample( batch_data, sample) return batch_finished def _finalise_minibatch(self, batch_data: Dict[str, Any], is_train: bool) -> Dict[tf.Tensor, Any]: minibatch = super()._finalise_minibatch(batch_data, is_train) self.__type_metric._finalise_minibatch(batch_data, is_train, minibatch) self.__type_classification._finalise_minibatch(batch_data, is_train, minibatch) return minibatch def create_index(self, data_paths: List[RichPath]): self.__type_metric.create_index(data_paths, self.metadata) # ------- These are the bits that we only need for test-time: def _encode_one_test_sample( self, sample_data_dict: Dict[tf.Tensor, Any] ) -> Tuple[tf.Tensor, Optional[tf.Tensor]]: return (self.sess.run(self.ops['target_representations'], feed_dict=sample_data_dict), None) def annotate_single(self, raw_sample: Dict[str, Any], loaded_test_sample: Dict[str, Any], provenance: str): return self.__type_metric.annotate_single(raw_sample, loaded_test_sample, provenance, self.metadata)