diff --git a/examples/finetune_integration.py b/examples/finetune_integration.py index 54ffc0e06..7386d120d 100644 --- a/examples/finetune_integration.py +++ b/examples/finetune_integration.py @@ -22,17 +22,11 @@ from torch.nn import functional as F from torch.utils.data import Dataset, DataLoader from sklearn.model_selection import train_test_split -from torchtext.vocab import Vocab -from torchtext._torchtext import ( - Vocab as VocabPybind, -) - -from scgpt.tokenizer.gene_tokenizer import GeneVocab sys.path.append("../") import scgpt as scg from scgpt.model import TransformerModel, AdversarialDiscriminator -from scgpt.tokenizer import tokenize_and_pad_batch, random_mask_value +from scgpt.tokenizer import GeneVocab, tokenize_and_pad_batch, random_mask_value from scgpt.loss import ( masked_mse_loss, masked_relative_error, @@ -222,8 +216,9 @@ # %% if config.load_model is None: - vocab = Vocab( - VocabPybind(genes + special_tokens, None) + tokens = genes + special_tokens + vocab = GeneVocab.from_dict( + {token: index for index, token in enumerate(tokens)} ) # bidirectional lookup [gene <-> int] vocab.set_default_index(vocab[""]) gene_ids = np.array(vocab(genes), dtype=int) @@ -795,4 +790,4 @@ def eval_testdata( run.finish() wandb.finish() -gc.collect() \ No newline at end of file +gc.collect() diff --git a/scgpt/tokenizer/gene_tokenizer.py b/scgpt/tokenizer/gene_tokenizer.py index 9227dd4ae..8898fdaa4 100644 --- a/scgpt/tokenizer/gene_tokenizer.py +++ b/scgpt/tokenizer/gene_tokenizer.py @@ -143,12 +143,24 @@ def from_dict( Args: token2idx (Dict[str, int]): Dictionary mapping tokens to indices. """ - # initiate an empty vocabulary first - _vocab = cls([], default_token=None) + for token, index in token2idx.items(): + if not isinstance(index, int) or isinstance(index, bool): + raise TypeError( + f"Vocabulary index for {token!r} must be an integer, " + f"got {type(index).__name__}." + ) + + ordered_items = sorted(token2idx.items(), key=lambda item: item[1]) + indices = [index for _, index in ordered_items] + if indices != list(range(len(ordered_items))): + raise ValueError( + "GeneVocab requires unique, consecutive indices starting at 0." + ) - # add the tokens to the vocabulary, GeneVocab requires consecutive indices - for t, i in sorted(token2idx.items(), key=lambda x: x[1]): - _vocab.insert_token(t, i) + # Initialize all tokens at once. Repeated insert_token calls rebuild the + # complete token-to-index mapping in the pure-Python backend. + _vocab = cls([], default_token=None) + _vocab._init_from_tokens([token for token, _ in ordered_items]) if default_token is not None and default_token in _vocab: _vocab.set_default_token(default_token) diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py index 545eb9f8a..179a70550 100644 --- a/tests/test_tokenizer.py +++ b/tests/test_tokenizer.py @@ -32,10 +32,31 @@ def test_gene_vocab(): def test_gene_vocab_from_dict(): - gene_vocab = GeneVocab.from_dict({"a": 0, "b": 1, "c": 2}) - assert len(gene_vocab) == 3 + gene_vocab = GeneVocab.from_dict({"c": 2, "": 3, "a": 0, "b": 1}) + assert len(gene_vocab) == 4 + assert gene_vocab.get_itos() == ["a", "b", "c", ""] assert gene_vocab["a"] == 0 assert gene_vocab["c"] == 2 + assert gene_vocab.get_default_index() == 3 + + +@pytest.mark.parametrize( + "token2idx", + [ + {"a": 0, "b": 2}, + {"a": 0, "b": 0}, + {"a": -1, "b": 0}, + ], +) +def test_gene_vocab_from_dict_rejects_invalid_index_sequence(token2idx): + with pytest.raises(ValueError, match="unique, consecutive indices"): + GeneVocab.from_dict(token2idx) + + +@pytest.mark.parametrize("invalid_index", [0.0, True, "0"]) +def test_gene_vocab_from_dict_rejects_non_integer_indices(invalid_index): + with pytest.raises(TypeError, match="must be an integer"): + GeneVocab.from_dict({"a": invalid_index}) def test_gene_vocab_from_file(): @@ -98,7 +119,7 @@ def test_builtin_vocab_append_insert(): v.insert_token("z", 0) assert v["z"] == 0 - assert v["a"] == 1 # shifted + assert v["a"] == 1 # shifted assert len(v) == 4 @@ -144,4 +165,3 @@ def test_gene_vocab_init_from_torchtext(): assert len(gv) == 3 for tok in tt_vocab.get_itos(): assert gv[tok] == tt_vocab[tok] -