Region2Vec Module

Region2VecExModel

class Region2VecExModel(model_path: str = None, tokenizer: Tokenizer = None, device: str = None, pooling_method: POOLING_TYPES = 'mean', **kwargs)

Bases: ExModel

Initialize Region2VecExModel.

Args: model_path (str): Path to the pre-trained model on huggingface. tokenizer (Tokenizer): Optional tokenizer to use. device (str): Device to use for the model. pooling_method (POOLING_TYPES): Pooling method to use. **kwargs: Additional keyword arguments to pass to the model.

Attributes

• model_path (str)
• tokenizer (Tokenizer)
• trained (bool)
• pooling_method

Properties

model

@property
def model

Get the core Region2Vec model.

Returns: Region2Vec: The core model.

Class methods

from_pretrained

@classmethod
def from_pretrained(path_to_files: str, model_file_name: str = MODEL_FILE_NAME, universe_file_name: str = UNIVERSE_FILE_NAME, config_file_name: str = CONFIG_FILE_NAME) -> Region2VecExModel

Load the model from a set of files that were exported using the export function.

Args: path_to_files (str): Path to the directory containing the files. model_file_name (str): Name of the model file. universe_file_name (str): Name of the universe file. config_file_name (str): Name of the config file.

Returns: Region2VecExModel: The loaded model instance.

Methods

add_tokenizer

def add_tokenizer(tokenizer: Tokenizer, **kwargs)

Add a tokenizer to the model.

This should be used when the model is not initialized with a tokenizer.

Args: tokenizer (Tokenizer): Tokenizer to add to the model. **kwargs: Additional keyword arguments to pass to the model.

Raises: RuntimeError: If the model is already initialized.

train

def train(dataset: Region2VecDataset, window_size: int = DEFAULT_WINDOW_SIZE, epochs: int = DEFAULT_EPOCHS, min_count: int = DEFAULT_MIN_COUNT, num_cpus: int = 1, seed: int = 42, save_checkpoint_path: str = None, gensim_params: dict = {}, load_from_checkpoint: str = None) -> bool

Train the model.

Args: dataset (Region2VecDataset): Dataset to train on. window_size (int): Window size for the model. epochs (int): Number of epochs to train for. min_count (int): Minimum count for a region to be included in the vocabulary. num_cpus (int): Number of cpus to use for training. seed (int): Seed to use for training. save_checkpoint_path (str): Path to save the model checkpoints to. gensim_params (dict): Additional parameters to pass to the gensim model. load_from_checkpoint (str): Path to a checkpoint to load from.

Returns: bool: True if training was successful.

Raises: RuntimeError: If model has not been initialized.

export

def export(path: str, checkpoint_file: str = MODEL_FILE_NAME, universe_file: str = UNIVERSE_FILE_NAME, config_file: str = CONFIG_FILE_NAME)

Export the model for huggingface upload.

This exports the model weights and the vocabulary in a format that can be directly uploaded to huggingface.

Args: path (str): Path to export the model to. checkpoint_file (str): Name of the checkpoint file. universe_file (str): Name of the universe file. config_file (str): Name of the config file.

encode

def encode(regions: Union[str, Region, Sequence[Region], RegionSet, GRegionSet], pooling: POOLING_TYPES = None, batch_size: Optional[int] = 64) -> np.ndarray

Vectorise one or many regions.

Args: regions (Union[str, Region, Sequence[Region], RegionSet, GRegionSet]): Region(s) to encode. pooling (POOLING_TYPES): "mean" or "max" token-pooling. batch_size (Optional[int]): How many regions to pad/encode at once. None or 0 processes all in one go.

Returns: np.ndarray: The encoded region vectors.

Raises: TypeError: If regions are not valid types. ValueError: If pooling method is invalid.

Word2Vec

class Word2Vec(vocab_size: int, embedding_dim: int = DEFAULT_EMBEDDING_DIM, **kwargs)

Bases: nn.Module

Word2Vec model.

Attributes

• vocab_size
• embedding_dim
• padding_idx
• projection

Methods

forward

def forward(x: torch.Tensor) -> torch.Tensor

Region2Vec

class Region2Vec(vocab_size: int, embedding_dim: int = DEFAULT_EMBEDDING_DIM, **kwargs)

Bases: Word2Vec

MeanPooling

class MeanPooling()

Bases: nn.Module

Methods

forward

def forward(x: torch.Tensor) -> torch.Tensor

MaxPooling

class MaxPooling()

Bases: nn.Module

Methods

forward

def forward(x: torch.Tensor) -> torch.Tensor

RegionSet2Vec

class RegionSet2Vec(region2vec: Region2Vec, pooling: POOLING_TYPES = 'mean')

Bases: nn.Module

Initialize the RegionSet2Vec model.

RegionSet2Vec is a wrapper around the Region2Vec model that allows pooling over a set of regions. This is useful for classification tasks where the input is a set of regions, such as classifying a cell type based on the set of regions that are accessible.

Args: region2vec (Union[Region2Vec, str]): Either a Region2Vec instance or a path to a huggingface model. pooling (POOLING_TYPES): The pooling type to use. Either "mean" or "max".

Raises: ValueError: If pooling type is invalid.

Attributes

• region2vec (Region2Vec)
• pooling

Methods

forward

def forward(x) -> torch.Tensor

mean_pooling

def mean_pooling(region_vectors: List[Union[np.ndarray, None]]) -> np.ndarray

Mean pooling of region vectors.

The function first filters out None values and then computes the mean of the remaining vectors.

Args: region_vectors (List[Union[np.ndarray, None]]): Region vectors to pool.

Returns: np.ndarray: The mean pooled vector.

max_pooling

def max_pooling(region_vectors: List[Union[np.ndarray, None]]) -> np.ndarray

Max pooling of region vectors.

The function first filters out None values and then computes the max of the remaining vectors.

Args: region_vectors (List[Union[np.ndarray, None]]): Region vectors to pool.

Returns: np.ndarray: The max pooled vector.

set_log_path

def set_log_path(path)

Timer

class Timer()

Records the running time.

Uses Timer.s() or Timer() to record the start time. Then, calls Timer.t() to get the elapsed time in seconds.

Initialize a Timer object and start the timer.

Attributes

• v

Methods

s

def s()

Restart the timer.

t

def t()

Give the elapsed time.

Returns: float: The elapsed time in seconds.

time_str

def time_str(t: float) -> str

Converts time in float to a readable format.

Converts time in float to hours, minutes, or seconds based on the value of t.

Args: t (float): Time in seconds.

Returns: str: Time in readable time.

timed_response

def timed_response(prompt: str, wait_time: int, default: str)

Prints prompt and waits for response.

Args: prompt (str): The question asks for a response. wait_time (int): The number of seconds for waiting. default (str): If no response received, uses default as the response.

Returns: str: a response given by the user or the default one.

log

def log(obj: str, filename: str = 'log.txt') -> None

Adds information in obj to a file specified by filename.

Adds information in obj to a file (default: log.txt) and prints obj.

Args: obj (str): A string. filename (str, optional): The log file name. Defaults to "log.txt".

lr_scheduler

class lr_scheduler(init_lr: float, end_lr: float, epochs: int, lr_info: Dict[str, Union[int, float, list]], mode: str = 'linear')

Changes the learning rate.

Changes the learning rate using the mode of linear or milestones. If mode = "linear", then the learning rate linearly decreases after certain epochs. If mode = "milestones", then the learning rate decreases at specified epochs.

Initialize the learning rate scheduler.

Args: init_lr (float): The initial learning rate. end_lr (float): The last learning rate. epochs (int): The number of training epochs. lr_info (Dict[str, Union[int, float, list]]): Dictionary storing information for learning rate scheduling. mode (str): The mode of learning rate scheduling. Defaults to "linear".

Attributes

• lr
• end_lr
• init_lr
• mode
• epochs
• lr_info
• count
• freq

Methods

step

def step()

Update the learning rate.

Returns: float: Current learning rate.

ensure_dir

def ensure_dir(folder: str, default: str = 'y') -> None

Make sure the folder exists.

If the folder exists, asks the user to keep [n] or delete [y] it. If no response received after 5 secs, deletes the folder and creates a new one.

Args: folder (str): The folder to be created. default (str): Choose whether to delete [y] or keep [n] the folder. Defaults to "y".

LearningRateScheduler

class LearningRateScheduler(init_lr: float = DEFAULT_INIT_LR, min_lr: float = DEFAULT_MIN_LR, type: LR_TYPES = 'exponential', decay: float = None, n_epochs: int = None)

Simple class to track learning rates of the training procedure.

Based off of: https://machinelearningmastery.com/using-learning-rate-schedules-deep-learning-models-python-keras/

Initialize the learning rate scheduler.

Args: init_lr (float): The initial learning rate. min_lr (float): The minimum learning rate. type (LR_TYPES): The type of learning rate schedule to use. Must be one of ['linear', 'exponential']. decay (float): The decay rate to use. If None, this will be calculated from init_lr and n_epochs. n_epochs (int): The number of epochs to train for. Only used if decay is None.

Attributes

• init_lr
• min_lr
• n_epochs
• type
• decay

Methods

update

def update()

get_lr

def get_lr()

shuffle_documents

def shuffle_documents(documents: List[List[any]], n_shuffles: int = 1, threads: int = None) -> List[List[any]]

Shuffle around the genomic regions for each cell to generate a "context".

Args: documents (List[List[any]]): The document list to shuffle. n_shuffles (int): The number of shuffles to conduct. threads (int): Optional number of threads to use for shuffling.

Returns: List[List[any]]: The shuffled documents.

export_region2vec_model

def export_region2vec_model(model: torch.nn.Module, tokenizer: Tokenizer, path: str, checkpoint_file: str = MODEL_FILE_NAME, universe_file: str = UNIVERSE_FILE_NAME, config_file: str = CONFIG_FILE_NAME, **kwargs: Dict[str, any])

Export the region2vec model to a folder.

Args: model (torch.nn.Module): The model to export. tokenizer (Tokenizer): The tokenizer to export. path (str): The path to export the model to. checkpoint_file (str): The name of the checkpoint file to export. universe_file (str): The name of the universe file to export. config_file (str): The name of the config file to export. **kwargs: Any additional arguments to pass to the config file.

load_local_region2vec_model

def load_local_region2vec_model(model_path: str, config_path: str, **kwargs) -> Tuple[Region2Vec, dict]

Load a region2vec model from a local directory.

Args: model_path (str): The path to the model checkpoint file. config_path (str): The path to the model config file. **kwargs: Include id of padding token.

Returns: Tuple[Region2Vec, dict]: The loaded model and its config.

Region2VecDataset

class Region2VecDataset(path: str, shuffle: bool = True, convert_to_str: bool = False)

Initialize a Region2VecDataset.

The regions are stored in a parquet file, with one document (cell, BED file, etc) per row.

Args: path (str): Path to the parquet file containing the tokens. shuffle (bool): Whether to shuffle the tokens in each document. convert_to_str (bool): Whether to convert the tokens to strings.

Attributes

• table
• data
• shuffle
• convert_to_str

train_region2vec_model

def train_region2vec_model(dataset: Region2VecDataset, embedding_dim: int = DEFAULT_EMBEDDING_DIM, window_size: int = DEFAULT_WINDOW_SIZE, epochs: int = DEFAULT_EPOCHS, min_count: int = DEFAULT_MIN_COUNT, num_cpus: int = 1, seed: int = 42, save_checkpoint_path: str = None, gensim_params: dict = {}, load_from_checkpoint: str = None) -> GensimWord2Vec

Train a gensim Word2Vec model on the given dataset.

Args: dataset (Region2VecDataset): Data to train on. This is a dataset of tokens. embedding_dim (int): Embedding dimension for the model. window_size (int): Window size for the model. epochs (int): Number of epochs to train for. min_count (int): Minimum count for a region to be included in the vocabulary. num_cpus (int): Number of cpus to use for training. seed (int): Seed to use for training. save_checkpoint_path (str): Path to save the model checkpoints to. gensim_params (dict): Additional parameters to pass to the gensim model. load_from_checkpoint (str): Path to a checkpoint to load from.

Returns: GensimWord2Vec: The gensim model that was trained.