biotransformers.wrappers.language_model#

This script defines a generic template class for any language model. Both ESM and Rostlab language models should implement this class.

Module Contents#

Classes#

LanguageModel

Class that implements a language model.
class biotransformers.wrappers.language_model.LanguageModel(model_dir: str, device)#

Bases: abc.ABC

Class that implements a language model.

property model_id(self) → str#

Model ID, as specified in the model directory

property clean_model_id(self) → str#

Clean model ID (in case the model directory is not)

property model_vocabulary(self) → List[str]#

Returns the whole vocabulary list

property vocab_size(self) → int#

Returns the whole vocabulary size

property mask_token(self) → str#

Representation of the mask token (as a string)

property pad_token(self) → str#

Representation of the pad token (as a string)

property begin_token(self) → str#

Representation of the beginning of sentence token (as a string)

property end_token(self) → str#

Representation of the end of sentence token (as a string).

property does_end_token_exist(self) → bool#

Returns true if a end of sequence token exists

property token_to_id(self)#

Returns a function which maps tokens to IDs

property embeddings_size(self) → int#

Returns size of the embeddings

abstract process_sequences_and_tokens(self, sequences_list: List[str]) → Dict[str, torch.Tensor]#

Function to transform tokens string to IDs; it depends on the model used

property model(self) → torch.nn.Module#

Return torch model.

abstract set_model(self, model: torch.nn.Module)#

Set torch model.

abstract model_pass(self, model_inputs: Dict[str, torch.tensor], batch_size: int, silent: bool = False, pba: ray.actor.ActorHandle = None) → Tuple[torch.Tensor, torch.Tensor]#

Function which computes logits and embeddings based on a dict of sequences tensors, a provided batch size and an inference configuration. The output is obtained by computing a forward pass through the model (“forward inference”)

Parameters

  • model_inputs (Dict[str, torch.tensor]) – [description]

  • batch_size (int) – size of the batch

  • silent – display or not progress bar

  • pba – tqdm progress bar for ray actor

Returns

  • logits [num_seqs, max_len_seqs, vocab_size]

  • embeddings [num_seqs, max_len_seqs+1, embedding_size]

Return type

Tuple[torch.tensor, torch.tensor]

abstract get_alphabet_dataloader(self)#

Define an alphabet mapping for common method between protbert and ESM