Improve gene vocabulary loading and torchtext compatibility

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["<pad>"])
 gene_ids = np.array(vocab(genes), dtype=int)
@@ -795,4 +790,4 @@ def eval_testdata(
 
 run.finish()
 wandb.finish()
-gc.collect()
+gc.collect()

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)

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, "<pad>": 3, "a": 0, "b": 1})
+    assert len(gene_vocab) == 4
+    assert gene_vocab.get_itos() == ["a", "b", "c", "<pad>"]
     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]
-