Upload tokenizer

special_tokens_map.json

@@ -0,0 +1 @@
+{}

tokenizer.py

@@ -0,0 +1,267 @@
+import json
+import os
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import PreTrainedTokenizer
+
+
+class BinnedOmicTokenizer(PreTrainedTokenizer):
+    def __init__(
+        self,
+        n_expressions_bins: int = 64,
+        min_omic_value: float = 0.0,
+        max_omic_value: float = 1.0,
+        use_max_normalization: bool = True,
+        normalization_factor: float = 1.0,
+        prepend_cls_token: bool = False,
+        fixed_sequence_length: Optional[int] = None,
+        unpadded_length: Optional[int] = None,
+        **kwargs,
+    ):
+        bin_tokens = [str(i) for i in range(n_expressions_bins)]
+        special_tokens = ["<pad>", "<mask>", "<cls>"]
+
+        vocab = {tok: i for i, tok in enumerate(bin_tokens)}
+        offset = len(vocab)
+        for i, tok in enumerate(special_tokens):
+            vocab[tok] = offset + i
+
+        ids_to_tokens = {i: tok for tok, i in vocab.items()}
+
+        self.vocab = vocab
+        self.ids_to_tokens = ids_to_tokens
+
+        self.n_expressions_bins = n_expressions_bins
+        self.min_omic_value = min_omic_value
+        self.max_omic_value = max_omic_value
+        self.use_max_normalization = use_max_normalization
+        self.normalization_factor = normalization_factor
+        self.prepend_cls_token = prepend_cls_token
+        self.fixed_sequence_length = fixed_sequence_length
+        self.unpadded_length = unpadded_length
+
+        self.bin_edges = np.linspace(min_omic_value, max_omic_value, n_expressions_bins)
+
+        self.pad_token = "<pad>"
+        self.mask_token = "<mask>"
+        self.cls_token = "<cls>"
+
+        super().__init__(**kwargs)
+
+        self.add_special_tokens(
+            {
+                "pad_token": "<pad>",
+                "mask_token": "<mask>",
+                "cls_token": "<cls>",
+                "unk_token": "<pad>",
+            }
+        )
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.vocab.get(token, self.vocab[self.unk_token])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.ids_to_tokens.get(index, self.unk_token)
+
+    def get_vocab(self) -> dict:
+        return self.vocab
+
+    def _tokenize(self, text, **kwargs):
+        raise NotImplementedError("Use `encode` or `batch_encode_plus` methods.")
+
+    def decode(self, token_ids, **kwargs):
+        return [self._convert_id_to_token(i) for i in token_ids]
+
+    def encode(
+        self,
+        gene_expr: Union[np.ndarray, List[float]],
+        pad_to_fixed_length: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ) -> Union[List[int], torch.Tensor]:
+        gene_expr = np.array(gene_expr)
+
+        if self.use_max_normalization:
+            gene_expr = gene_expr / self.normalization_factor
+
+        token_ids = np.digitize(gene_expr, self.bin_edges).astype(int)
+        token_ids[gene_expr == 0.0] = 0
+
+        if self.prepend_cls_token:
+            token_ids = np.concatenate([[self.cls_token_id], token_ids])
+
+        if pad_to_fixed_length:
+            current_max_length = self.fixed_sequence_length or max_length
+            if current_max_length is None:
+                raise ValueError("fixed_sequence_length or max_length must be set.")
+            pad_len = current_max_length - len(token_ids)
+            if pad_len > 0:
+                token_ids = np.concatenate([token_ids, [self.pad_token_id] * pad_len])
+            else:
+                token_ids = token_ids[:current_max_length]
+
+        if return_tensors == "pt":
+            return torch.tensor(token_ids).unsqueeze(0)
+        return token_ids.tolist()  # type: ignore
+
+    def batch_encode_plus(
+        self,
+        batch_gene_expr: Union[np.ndarray, List[np.ndarray]],
+        pad_to_fixed_length: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ):
+        if isinstance(batch_gene_expr, list):
+            batch_gene_expr = np.array(batch_gene_expr)
+
+        encoded = [
+            self.encode(
+                gene_expr,
+                pad_to_fixed_length=pad_to_fixed_length,
+                max_length=max_length,
+                return_tensors=None,
+                **kwargs,
+            )
+            for gene_expr in batch_gene_expr
+        ]
+
+        encoded = np.array(encoded, dtype=np.int64)
+
+        if return_tensors == "pt":
+            return {"input_ids": torch.tensor(encoded)}
+        return {"input_ids": encoded}
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.vocab)
+
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ):
+        vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + "vocab.json",
+        )
+        with open(vocab_file, "w") as f:
+            json.dump(self.vocab, f)
+        return (vocab_file,)
+
+
+class MOJOTokenizer(PreTrainedTokenizer):
+    def __init__(
+        self,
+        n_expressions_bins: dict[str, int],
+        min_omic_value: dict[str, float],
+        max_omic_value: dict[str, float],
+        use_max_normalization: dict[str, bool],
+        normalization_factor: dict[str, float],
+        prepend_cls_token: bool,
+        fixed_sequence_length: int,
+        unpadded_length: int,
+        **kwargs,
+    ):
+        self.omics = n_expressions_bins.keys()
+        self.omic_tokenizers = {
+            omic: BinnedOmicTokenizer(
+                n_expressions_bins=n_expressions_bins[omic],
+                min_omic_value=min_omic_value[omic],
+                max_omic_value=max_omic_value[omic],
+                use_max_normalization=use_max_normalization[omic],
+                normalization_factor=normalization_factor[omic],
+                prepend_cls_token=prepend_cls_token,
+                fixed_sequence_length=fixed_sequence_length,
+                unpadded_length=unpadded_length,
+                **kwargs,
+            )
+            for omic in n_expressions_bins.keys()
+        }
+
+        self.vocab = {omic: self.omic_tokenizers[omic].vocab for omic in self.omics}
+        self.ids_to_tokens = {
+            omic: self.omic_tokenizers[omic].ids_to_tokens for omic in self.omics
+        }
+
+        super().__init__(**kwargs)
+
+    def _convert_token_to_id(self, token: dict[str, str]) -> dict[str, int]:
+        return {
+            omic: self.vocab[omic].get(token[omic], self.vocab[omic][self.unk_token])
+            for omic in token
+        }
+
+    def _convert_id_to_token(self, index: dict[str, int]) -> dict[str, str]:
+        return {
+            omic: self.omic_tokenizers[omic]._convert_id_to_token(index[omic])
+            for omic in index
+        }
+
+    def get_vocab(self) -> dict:
+        return self.vocab
+
+    def _tokenize(self, text, **kwargs):
+        raise NotImplementedError("Use `encode` or `batch_encode_plus` methods.")
+
+    def decode(self, token_ids: dict[str, list[int]], **kwargs):
+        return {
+            omic: self.omic_tokenizers[omic].decode(token_ids[omic])
+            for omic in token_ids
+        }
+
+    def encode(
+        self,
+        omic_array: Union[dict[str, np.ndarray], dict[str, List[float]]],
+        pad_to_fixed_length: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ) -> Union[dict[str, List[int]], dict[str, torch.Tensor]]:
+        return {
+            omic: self.omic_tokenizers[omic].encode(
+                omic_array[omic],
+                pad_to_fixed_length=pad_to_fixed_length,
+                max_length=max_length,
+                return_tensors=return_tensors,
+            )
+            for omic in omic_array
+        }
+
+    def batch_encode_plus(
+        self,
+        batch_omic_array: Union[dict[str, np.ndarray], dict[str, List[np.ndarray]]],
+        pad_to_fixed_length: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ):
+        return {
+            omic: self.omic_tokenizers[omic].batch_encode_plus(
+                batch_omic_array[omic],
+                pad_to_fixed_length=pad_to_fixed_length,
+                max_length=max_length,
+                return_tensors=return_tensors,
+            )
+            for omic in batch_omic_array
+        }
+
+    @property
+    def vocab_size(self) -> int:
+        return sum(len(self.vocab[omic]) for omic in self.vocab)
+
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ):
+        vocab_files = []
+        for omic in self.omics:
+            vocab_file = os.path.join(
+                save_directory,
+                (filename_prefix + "-" if filename_prefix else "")
+                + f"vocab_{omic}.json",
+            )
+            with open(vocab_file, "w") as f:
+                json.dump(self.vocab[omic], f)
+            vocab_files.append(vocab_file)
+        return tuple(vocab_files)

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "added_tokens_decoder": {},
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenizer.MOJOTokenizer",
+      null
+    ]
+  },
+  "clean_up_tokenization_spaces": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "MOJOTokenizer"
+}

vocab_methylation.json

@@ -0,0 +1 @@
+{"0": 0, "1": 1, "2": 2, "3": 3, "4": 4, "5": 5, "6": 6, "7": 7, "8": 8, "9": 9, "10": 10, "11": 11, "12": 12, "13": 13, "14": 14, "15": 15, "16": 16, "17": 17, "18": 18, "19": 19, "20": 20, "21": 21, "22": 22, "23": 23, "24": 24, "25": 25, "26": 26, "27": 27, "28": 28, "29": 29, "30": 30, "31": 31, "32": 32, "33": 33, "34": 34, "35": 35, "36": 36, "37": 37, "38": 38, "39": 39, "40": 40, "41": 41, "42": 42, "43": 43, "44": 44, "45": 45, "46": 46, "47": 47, "48": 48, "49": 49, "50": 50, "51": 51, "52": 52, "53": 53, "54": 54, "55": 55, "56": 56, "57": 57, "58": 58, "59": 59, "60": 60, "61": 61, "62": 62, "63": 63, "<pad>": 64, "<mask>": 65, "<cls>": 66}

vocab_rnaseq.json

@@ -0,0 +1 @@
+{"0": 0, "1": 1, "2": 2, "3": 3, "4": 4, "5": 5, "6": 6, "7": 7, "8": 8, "9": 9, "10": 10, "11": 11, "12": 12, "13": 13, "14": 14, "15": 15, "16": 16, "17": 17, "18": 18, "19": 19, "20": 20, "21": 21, "22": 22, "23": 23, "24": 24, "25": 25, "26": 26, "27": 27, "28": 28, "29": 29, "30": 30, "31": 31, "32": 32, "33": 33, "34": 34, "35": 35, "36": 36, "37": 37, "38": 38, "39": 39, "40": 40, "41": 41, "42": 42, "43": 43, "44": 44, "45": 45, "46": 46, "47": 47, "48": 48, "49": 49, "50": 50, "51": 51, "52": 52, "53": 53, "54": 54, "55": 55, "56": 56, "57": 57, "58": 58, "59": 59, "60": 60, "61": 61, "62": 62, "63": 63, "<pad>": 64, "<mask>": 65, "<cls>": 66}