masskit_ai.mol.small.models package¶

Submodules¶

masskit_ai.mol.small.models.algos module¶

masskit_ai.mol.small.models.gf module¶

class masskit_ai.mol.small.models.gf.BaseFairseqModel(*args: Any, **kwargs: Any)¶

Bases: Module

Base class for fairseq models.

classmethod build_model(args, task)¶: Build a new model instance.

extract_features(*args, **kwargs)¶: Similar to forward but only return features.

get_normalized_probs(net_output: Tuple[torch.Tensor, Optional[Dict[str, List[Optional[torch.Tensor]]]]], log_probs: bool, sample: Optional[Dict[str, torch.Tensor]] = None)¶: Get normalized probabilities (or log probs) from a net’s output.

get_normalized_probs_scriptable(net_output: Tuple[torch.Tensor, Optional[Dict[str, List[Optional[torch.Tensor]]]]], log_probs: bool, sample: Optional[Dict[str, torch.Tensor]] = None)¶: Scriptable helper function for get_normalized_probs in ~BaseFairseqModel

get_targets(sample, net_output)¶: Get targets from either the sample or the net’s output.

classmethod hub_models()¶

make_generation_fast_(**kwargs)¶: Legacy entry point to optimize model for faster generation.

max_positions()¶: Maximum length supported by the model.

prepare_for_inference_(cfg: omegaconf.DictConfig)¶: Prepare model for inference.

prepare_for_onnx_export_(**kwargs)¶: Make model exportable via ONNX trace.

set_num_updates(num_updates)¶: State from trainer to pass along to model at every update.

class masskit_ai.mol.small.models.gf.FairseqDropout(*args: Any, **kwargs: Any)¶

Bases: Module

forward(x, inplace: bool = False)¶

make_generation_fast_(name: str, retain_dropout: bool = False, retain_dropout_modules: Optional[List[str]] = None, **kwargs)¶

class masskit_ai.mol.small.models.gf.FairseqEncoder(*args: Any, **kwargs: Any)¶

Bases: Module

Base class for encoders.

forward(src_tokens, src_lengths=None, **kwargs)¶

Args:

src_tokens (LongTensor): tokens in the source language of shape: (batch, src_len)
src_lengths (LongTensor): lengths of each source sentence of shape: (batch)

forward_non_torchscript(net_input: Dict[str, torch.Tensor])¶

forward_torchscript(net_input: Dict[str, torch.Tensor])¶

A TorchScript-compatible version of forward.

Encoders which use additional arguments may want to override this method for TorchScript compatibility.

max_positions()¶: Maximum input length supported by the encoder.

reorder_encoder_out(encoder_out, new_order)¶

Reorder encoder output according to new_order.

Args:: encoder_out: output from the forward() method new_order (LongTensor): desired order
Returns:: encoder_out rearranged according to new_order

set_num_updates(num_updates)¶: State from trainer to pass along to model at every update.

upgrade_state_dict_named(state_dict, name)¶: Upgrade old state dicts to work with newer code.

class masskit_ai.mol.small.models.gf.FairseqEncoderModel(*args: Any, **kwargs: Any)¶

Bases: BaseFairseqModel

Base class for encoder-only models.

Args:: encoder (FairseqEncoder): the encoder

forward(src_tokens, src_lengths, **kwargs)¶

Run the forward pass for a encoder-only model.

Feeds a batch of tokens through the encoder to generate features.

Args:: src_tokens (LongTensor): input tokens of shape (batch, src_len) src_lengths (LongTensor): source sentence lengths of shape (batch)
Returns:: the encoder’s output, typically of shape (batch, src_len, features)

get_normalized_probs(net_output, log_probs, sample=None)¶: Get normalized probabilities (or log probs) from a net’s output.

max_positions()¶: Maximum length supported by the model.

class masskit_ai.mol.small.models.gf.GraphAttnBias(*args: Any, **kwargs: Any)¶

Bases: Module

Compute attention bias for each head.

forward(batched_data)¶

class masskit_ai.mol.small.models.gf.GraphNodeFeature(*args: Any, **kwargs: Any)¶

Bases: Module

Compute node features for each node in the graph.

forward(batched_data)¶

class masskit_ai.mol.small.models.gf.GraphormerEncoder(*args: Any, **kwargs: Any)¶

Bases: FairseqEncoder

forward(batched_data, perturb=None, masked_tokens=None, **unused)¶

Args:

src_tokens (LongTensor): tokens in the source language of shape: (batch, src_len)
src_lengths (LongTensor): lengths of each source sentence of shape: (batch)

max_nodes()¶: Maximum output length supported by the encoder.

reset_output_layer_parameters()¶

upgrade_state_dict_named(state_dict, name)¶: Upgrade old state dicts to work with newer code.

class masskit_ai.mol.small.models.gf.GraphormerGraphEncoder(*args: Any, **kwargs: Any)¶

Bases: Module

build_graphormer_graph_encoder_layer(embedding_dim, ffn_embedding_dim, num_attention_heads, dropout, attention_dropout, activation_dropout, activation_fn, export, q_noise, qn_block_size, pre_layernorm)¶

forward(batched_data, perturb=None, last_state_only: bool = False, token_embeddings: Optional[torch.Tensor] = None, attn_mask: Optional[torch.Tensor] = None) → Tuple[torch.Tensor, torch.Tensor]¶

class masskit_ai.mol.small.models.gf.GraphormerGraphEncoderLayer(*args: Any, **kwargs: Any)¶

Bases: Module

build_fc1(input_dim, output_dim, q_noise, qn_block_size)¶

build_fc2(input_dim, output_dim, q_noise, qn_block_size)¶

build_self_attention(embed_dim, num_attention_heads, dropout, self_attention, q_noise, qn_block_size)¶

forward(x: torch.Tensor, self_attn_bias: Optional[torch.Tensor] = None, self_attn_mask: Optional[torch.Tensor] = None, self_attn_padding_mask: Optional[torch.Tensor] = None)¶: LayerNorm is applied either before or after the self-attention/ffn modules similar to the original Transformer implementation.

class masskit_ai.mol.small.models.gf.GraphormerModel(*args: Any, **kwargs: Any)¶

Bases: FairseqEncoderModel

classmethod build_model(config)¶: Build a new model instance.

forward(batched_data, **kwargs)¶

Run the forward pass for a encoder-only model.

Feeds a batch of tokens through the encoder to generate features.

Args:: src_tokens (LongTensor): input tokens of shape (batch, src_len) src_lengths (LongTensor): source sentence lengths of shape (batch)
Returns:: the encoder’s output, typically of shape (batch, src_len, features)

max_nodes()¶

class masskit_ai.mol.small.models.gf.Graphormer_slim(*args: Any, **kwargs: Any)¶

Bases: SpectrumModel

class for wrapping graphormer in a standard SpectrumModel

forward(inp)¶

total_params()¶

class masskit_ai.mol.small.models.gf.LayerDropModuleList(*args: Any, **kwargs: Any)¶

Bases: ModuleList

A LayerDrop implementation based on torch.nn.ModuleList.

We refresh the choice of which layers to drop every time we iterate over the LayerDropModuleList instance. During evaluation we always iterate over all layers.

Usage:

layers = LayerDropList(p=0.5, modules=[layer1, layer2, layer3])
for layer in layers:  # this might iterate over layers 1 and 3
    x = layer(x)
for layer in layers:  # this might iterate over all layers
    x = layer(x)
for layer in layers:  # this might not iterate over any layers
    x = layer(x)

Args:: p (float): probability of dropping out each layer modules (iterable, optional): an iterable of modules to add

masskit_ai.mol.small.models.gf.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, export=False)¶

class masskit_ai.mol.small.models.gf.MultiheadAttention(*args: Any, **kwargs: Any)¶

Bases: Module

Multi-headed attention.

See “Attention Is All You Need” for more details.

apply_sparse_mask(attn_weights, tgt_len: int, src_len: int, bsz: int)¶

forward(query, key: Optional[torch.Tensor], value: Optional[torch.Tensor], attn_bias: Optional[torch.Tensor], key_padding_mask: Optional[torch.Tensor] = None, need_weights: bool = True, attn_mask: Optional[torch.Tensor] = None, before_softmax: bool = False, need_head_weights: bool = False) → Tuple[torch.Tensor, Optional[torch.Tensor]]¶

Input shape: Time x Batch x Channel

Args:

key_padding_mask (ByteTensor, optional): mask to exclude: keys that are pads, of shape (batch, src_len), where padding elements are indicated by 1s.
need_weights (bool, optional): return the attention weights,: averaged over heads (default: False).
attn_mask (ByteTensor, optional): typically used to: implement causal attention, where the mask prevents the attention from looking forward in time (default: None).
before_softmax (bool, optional): return the raw attention: weights and values before the attention softmax.
need_head_weights (bool, optional): return the attention: weights for each head. Implies need_weights. Default: return the average attention weights over all heads.

prepare_for_onnx_export_()¶

reset_parameters()¶

upgrade_state_dict_named(state_dict, name)¶

masskit_ai.mol.small.models.gf.base_architecture(args)¶

masskit_ai.mol.small.models.gf.check_type(module, expected_type)¶

masskit_ai.mol.small.models.gf.gelu(x: torch.Tensor) → torch.Tensor¶

masskit_ai.mol.small.models.gf.gelu_accurate(x)¶

masskit_ai.mol.small.models.gf.get_activation_fn(activation: str) → Callable¶: Returns the activation function corresponding to activation

masskit_ai.mol.small.models.gf.graphormer_slim_architecture(config)¶

masskit_ai.mol.small.models.gf.init_graphormer_params(module)¶: Initialize the weights specific to the Graphormer Model.

masskit_ai.mol.small.models.gf.init_params(module, n_layers)¶

masskit_ai.mol.small.models.gf.quant_noise(module, p, block_size)¶

Wraps modules and applies quantization noise to the weights for subsequent quantization with Iterative Product Quantization as described in “Training with Quantization Noise for Extreme Model Compression”

Args:

module: nn.Module
p: amount of Quantization Noise
block_size: size of the blocks for subsequent quantization with iPQ

Remarks:

Module weights must have the right sizes wrt the block size
Only Linear, Embedding and Conv2d modules are supported for the moment
For more detail on how to quantize by blocks with convolutional weights, see “And the Bit Goes Down: Revisiting the Quantization of Neural Networks”
We implement the simplest form of noise here as stated in the paper which consists in randomly dropping blocks

masskit_ai.mol.small.models.gf.relu_squared(x: torch.Tensor)¶

masskit_ai.mol.small.models.gf.safe_hasattr(obj, k)¶: Returns True if the given key exists and is not None.

masskit_ai.mol.small.models.gf.softmax(x, dim: int, onnx_trace: bool = False)¶

masskit_ai.mol.small.models.model_utils module¶

masskit_ai.mol.small.models.model_utils.compute_max_atoms(scope)¶: Given the scopes for a batch of molecules, computes max atoms.

masskit_ai.mol.small.models.model_utils.convert_to_2D(input, scope)¶

Convert back to 2D

Args:: input: A tensor of shape [batch size, max padding, # features] scope: A list of start/length indices for the molecules
Returns:: A matrix of size [# atoms, # features]

masskit_ai.mol.small.models.model_utils.convert_to_3D(input, scope, max_atoms, self_attn=True)¶

Converts the input to a 3D batch matrix

Args:: input: A tensor of shape [# atoms, # features] scope: A list of start/length indices for the molecules max_atoms: The maximum number of atoms for padding purposes
Returns:: A matrix of size [batch_size, max atoms, # features]

masskit_ai.mol.small.models.mol_features module¶

masskit_ai.mol.small.models.mol_features.bt_index_to_float(bt_index)¶

masskit_ai.mol.small.models.mol_features.get_atom_features(atom)¶: Given an atom object, returns a numpy array of features.

masskit_ai.mol.small.models.mol_features.get_bond_features(bond, bt_only=False)¶

Given an bond object, returns a numpy array of features.

bond can be None, in which case returns default features for a non-bond.

masskit_ai.mol.small.models.mol_features.get_bt_feature(bond_type)¶: Returns a one-hot vector representing the bond_type.

masskit_ai.mol.small.models.mol_features.get_bt_index(bond_type)¶

Returns the feature index for a particular bond type.

Args:

bond_type: Either a rdchem bond type object (can be None) or a float: representing the bond type

masskit_ai.mol.small.models.mol_features.get_path_bond_feature(bond)¶

Given a rdkit bond object, returns the bond features for that bond.

When the given input is none, returns a 0-vector

masskit_ai.mol.small.models.mol_features.onek_unk_encoding(x, set)¶: Returns a one-hot encoding of the given feature.

masskit_ai.mol.small.models.mol_graph module¶

class masskit_ai.mol.small.models.mol_graph.Atom(idx, rd_atom=None, is_dummy=False)¶

Bases: object

add_bond(bond)¶

class masskit_ai.mol.small.models.mol_graph.Bond(idx, out_atom_idx, in_atom_idx, rd_bond=None)¶: Bases: object

class masskit_ai.mol.small.models.mol_graph.MolGraph(rd_mols, path_input=None, path_mask=None)¶

Bases: object

get_atom_inputs(output_tensors=True)¶: Constructs only the atom inputs for the batch of molecules.

get_graph_inputs(output_tensors=True)¶

Constructs the graph inputs for the conv net.

Returns:: A tuple of tensors/numpy arrays that contains the input to the GCN.

get_n_atoms()¶

class masskit_ai.mol.small.models.mol_graph.Molecule(atoms, bonds)¶

Bases: object

get_bond(atom_1, atom_2)¶

masskit_ai.mol.small.models.path_utils module¶

masskit_ai.mol.small.models.path_utils.get_num_path_features(args)¶: Returns the number of path features for the model.

masskit_ai.mol.small.models.path_utils.get_path_atoms(atom_1, atom_2, paths_dict, pointer_dict, max_path_length, truncate=True, self_attn=False)¶

Given a pair of atom indices, returns the list of atoms on the path.

Parameters:

atom_1 – The start atom on the path.
atom_2 – The end atom on the path.
paths_dict – A mapping from atom pairs to paths.
pointer_dict – A mapping from atom pairs to truncated paths. The values of the dict is the atom ending the truncated path
max_path_length – The maximum path length to return.
truncate – Boolean determining whether or not paths above the max length should be truncated (returned as no path)

masskit_ai.mol.small.models.path_utils.get_path_features(rd_mol, path_atoms, path_length, max_path_length, p_embed=False)¶

Returns a feature array for the path.

Parameters:

rd_mol – The rdkit mol object, used to extract features.
path_atoms – A list of atoms in the path. If no path exist, empty array.
path_length – The length of the path.
max_path_length – The maximum length of the path considered.
p_embed – Whether or not to use position embedding.

masskit_ai.mol.small.models.path_utils.get_path_input(mols, shortest_paths, max_atoms, args, output_tensor=True)¶

masskit_ai.mol.small.models.path_utils.get_ring_features(ring_dict, atom_pair)¶

masskit_ai.mol.small.models.path_utils.get_value_attr(some_thing, some_key)¶

masskit_ai.mol.small.models.path_utils.merge_path_inputs(path_inputs, path_masks, max_atoms, args)¶: Merge path input matrices. Does not create CUDA tensors intentionally

masskit_ai.mol.small.models.path_utils.ordered_pair(a1, a2)¶

masskit_ai.mol.small.models.patn module¶

class masskit_ai.mol.small.models.patn.MolTransformer(*args: Any, **kwargs: Any)¶

Bases: Module

avg_attn(attn_probs, n_heads, batch_sz, max_atoms)¶

compute_attn_probs(attn_input, attn_mask, layer_idx, eps=1e-20)¶

forward(mol_graph)¶

get_attn_input(atom_h, path_input, max_atoms)¶

class masskit_ai.mol.small.models.patn.PropPredictor(*args: Any, **kwargs: Any)¶

Bases: SpectrumModel

aggregate_atom_h(atom_h, scope)¶

forward(input, output_attn=False)¶

masskit_ai.mol.small.models package¶

Submodules¶

masskit_ai.mol.small.models.algos module¶

masskit_ai.mol.small.models.gf module¶

masskit_ai.mol.small.models.model_utils module¶

masskit_ai.mol.small.models.mol_features module¶

masskit_ai.mol.small.models.mol_graph module¶

masskit_ai.mol.small.models.path_utils module¶

masskit_ai.mol.small.models.patn module¶

Module contents¶

Table of Contents

This Page