from typing import List, Dict, Optional, Union, Tuple import numpy as np from transformers import PreTrainedTokenizer from dataclasses import dataclass import torch import anndata as ad from scipy.sparse import issparse import numba import os import json # Token IDs must match exactly with the original implementation PAD_TOKEN = 0 MASK_TOKEN = 1 CLS_TOKEN = 2 # These mappings preserve the exact token IDs from the original implementation MODALITY_DICT = { 'dissociated': 3, 'spatial': 4, } SPECIES_DICT = { 'human': 5, 'Homo sapiens': 5, 'Mus musculus': 6, 'mouse': 6, } TECHNOLOGY_DICT = { "merfish": 7, "MERFISH": 7, "cosmx": 8, "NanoString digital spatial profiling": 8, "Xenium": 9, "10x 5' v2": 10, "10x 3' v3": 11, "10x 3' v2": 12, "10x 5' v1": 13, "10x 3' v1": 14, "10x 3' transcription profiling": 15, "10x transcription profiling": 15, "10x 5' transcription profiling": 16, "CITE-seq": 17, "Smart-seq v4": 18, } def sf_normalize(X: np.ndarray) -> np.ndarray: """Size factor normalize to 10k counts.""" X = X.copy() counts = np.array(X.sum(axis=1)) # avoid zero division error counts += counts == 0. # normalize to 10000 counts scaling_factor = 10000. / counts if issparse(X): from scipy.sparse import sparsefuncs sparsefuncs.inplace_row_scale(X, scaling_factor) else: np.multiply(X, scaling_factor.reshape((-1, 1)), out=X) return X @numba.jit(nopython=True, nogil=True) def _sub_tokenize_data(x: np.ndarray, max_seq_len: int = -1, aux_tokens: int = 30) -> np.ndarray: """Tokenize the input gene vector.""" scores_final = np.empty((x.shape[0], max_seq_len if max_seq_len > 0 else x.shape[1])) for i, cell in enumerate(x): nonzero_mask = np.nonzero(cell)[0] sorted_indices = nonzero_mask[np.argsort(-cell[nonzero_mask])][:max_seq_len] sorted_indices = sorted_indices + aux_tokens if max_seq_len: scores = np.zeros(max_seq_len, dtype=np.int32) else: scores = np.zeros_like(cell, dtype=np.int32) scores[:len(sorted_indices)] = sorted_indices.astype(np.int32) scores_final[i, :] = scores return scores_final class NicheformerTokenizer(PreTrainedTokenizer): """Tokenizer for Nicheformer that handles single-cell data.""" model_input_names = ["input_ids", "attention_mask"] vocab_files_names = {"vocab_file": "vocab.json"} modality_dict = MODALITY_DICT species_dict = SPECIES_DICT technology_dict = TECHNOLOGY_DICT def __init__( self, vocab_file=None, max_length: int = 1500, aux_tokens: int = 30, median_counts_per_gene: Optional[np.ndarray] = None, gene_names: Optional[List[str]] = None, **kwargs ): # Initialize base vocabulary self._vocabulary = { "[PAD]": PAD_TOKEN, "[MASK]": MASK_TOKEN, "[CLS]": CLS_TOKEN, } if vocab_file is not None: with open(vocab_file, 'r') as f: self._vocabulary.update(json.load(f)) else: # Add modality tokens for name, idx in self.modality_dict.items(): self._vocabulary[f"[MODALITY_{name}]"] = idx # Add species tokens for name, idx in self.species_dict.items(): if name in ["Homo sapiens", "Mus musculus"]: continue # Skip redundant names self._vocabulary[f"[SPECIES_{name}]"] = idx # Add technology tokens for name, idx in self.technology_dict.items(): if name in ["MERFISH", "10x transcription profiling"]: continue # Skip redundant names clean_name = name.lower().replace(" ", "_").replace("'", "_") self._vocabulary[f"[TECH_{clean_name}]"] = idx # Add gene tokens if provided if gene_names is not None: for i, gene in enumerate(gene_names): self._vocabulary[gene] = i + aux_tokens # Save vocabulary os.makedirs('to_hf', exist_ok=True) with open('to_hf/vocab.json', 'w') as f: json.dump(self._vocabulary, f, indent=4) super().__init__(**kwargs) self.max_length = max_length self.aux_tokens = aux_tokens self.median_counts_per_gene = median_counts_per_gene self.gene_names = gene_names # Set up special token mappings self._pad_token = "[PAD]" self._mask_token = "[MASK]" self._cls_token = "[CLS]" def get_vocab(self) -> Dict[str, int]: """Returns the vocabulary mapping.""" return self._vocabulary.copy() def _tokenize(self, text: str) -> List[str]: """Tokenize text input.""" # This tokenizer doesn't handle text input directly raise NotImplementedError("This tokenizer only works with gene expression data") def _convert_token_to_id(self, token: str) -> int: """Convert token to ID.""" # First check special token mappings if token in self.modality_dict: return self.modality_dict[token] if token in self.species_dict: return self.species_dict[token] if token in self.technology_dict: return self.technology_dict[token] # Then check vocabulary return self._vocabulary.get(token, self._vocabulary["[PAD]"]) def _convert_id_to_token(self, index: int) -> str: """Convert ID to token.""" # First check special token mappings for token, idx in self.modality_dict.items(): if idx == index: return token for token, idx in self.species_dict.items(): if idx == index: return token for token, idx in self.technology_dict.items(): if idx == index: return token # Then check vocabulary for token, idx in self._vocabulary.items(): if idx == index: return token return "[PAD]" def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]: """Save the vocabulary to a file.""" vocab_file = os.path.join( save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.json" ) with open(vocab_file, "w", encoding="utf-8") as f: json.dump(self._vocabulary, f, ensure_ascii=False) return (vocab_file,) def _tokenize_gene_expression(self, x: np.ndarray) -> np.ndarray: """Tokenize gene expression matrix. Args: x: Gene expression matrix (cells x genes) Returns: Tokenized matrix """ # Handle sparse input if issparse(x): x = x.toarray() # Normalize and scale x = np.nan_to_num(x) x = sf_normalize(x) if self.median_counts_per_gene is not None: median_counts = self.median_counts_per_gene.copy() median_counts += median_counts == 0 x = x / median_counts.reshape((1, -1)) # Convert to tokens tokens = _sub_tokenize_data(x, self.max_length, self.aux_tokens) return tokens.astype(np.int32) def __call__( self, adata: Optional[ad.AnnData] = None, gene_expression: Optional[Union[np.ndarray, List[float]]] = None, modality: Optional[str] = None, species: Optional[str] = None, technology: Optional[str] = None, **kwargs ) -> Dict[str, torch.Tensor]: """Convert inputs to model inputs. Args: adata: AnnData object gene_expression: Gene expression matrix if not using AnnData modality: Modality type species: Species type technology: Technology/assay type Returns: Dictionary with model inputs """ if adata is not None: # Get expression matrix if issparse(adata.X): x = adata.X.toarray() else: x = adata.X # Get metadata for each cell if not provided if modality is None and 'modality' in adata.obs: modality = adata.obs['modality'].values if species is None and 'specie' in adata.obs: species = adata.obs['specie'].values if technology is None and 'assay' in adata.obs: technology = adata.obs['assay'].values elif gene_expression is not None: x = np.array(gene_expression) if len(x.shape) == 1: x = x.reshape(1, -1) # For single gene expression input, convert scalar metadata to arrays if modality is not None: modality = np.array([modality]) if species is not None: species = np.array([species]) if technology is not None: technology = np.array([technology]) else: raise ValueError("Either adata or gene_expression must be provided") # Tokenize gene expression token_ids = self._tokenize_gene_expression(x) n_cells = token_ids.shape[0] # Add special tokens for each cell special_tokens = np.zeros((n_cells, 3), dtype=np.int32) # 3 for modality, species, technology special_token_mask = np.zeros((n_cells, 3), dtype=bool) # Track which tokens are actually present if modality is not None: special_tokens[:, 0] = [self.modality_dict.get(m, self._vocabulary["[PAD]"]) for m in modality] special_token_mask[:, 0] = True if species is not None: special_tokens[:, 1] = [self.species_dict.get(s, self._vocabulary["[PAD]"]) for s in species] special_token_mask[:, 1] = True if technology is not None: special_tokens[:, 2] = [self.technology_dict.get(t, self._vocabulary["[PAD]"]) for t in technology] special_token_mask[:, 2] = True # Only keep the special tokens that are present (have True in mask) special_tokens = special_tokens[:, special_token_mask[0]] if special_tokens.size > 0: token_ids = np.concatenate([special_tokens, token_ids[:, :(self.max_length - special_tokens.shape[1])]], axis=1) # Create attention mask attention_mask = (token_ids != self._vocabulary["[PAD]"]) return { "input_ids": torch.tensor(token_ids, dtype=torch.long), "attention_mask": torch.tensor(attention_mask) } def get_vocab_size(self) -> int: """Get vocabulary size.""" if self.gene_names is not None: return len(self.gene_names) + self.aux_tokens return max( max(self.modality_dict.values()), max(self.species_dict.values()), max(self.technology_dict.values()) ) + 1 def convert_tokens_to_string(self, tokens: List[str]) -> str: """Convert a sequence of tokens to a string. Not used for gene expression.""" raise NotImplementedError("This tokenizer only works with gene expression data") def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None) -> List[int]: """Build model inputs from a sequence by adding special tokens.""" # For gene expression data, special tokens are handled in __call__ return token_ids_0 def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False) -> List[int]: """Get list where entries are [1] if a token is [special] else [0].""" # Consider tokens < aux_tokens as special return [1 if token_id < self.aux_tokens else 0 for token_id in token_ids_0]