fix: add tests

python/mtc_token_healing/mtc_token_healing.pyi

@@ -1,10 +1,18 @@
+from typing import Sequence
+
 TokenId = int
 
 class BestChoice: ...
 class CountInfo: ...
 class InferRequest: ...
 class InferResponse: ...
 class Prediction: ...
-class VocabPrefixAutomaton: ...
+
+class VocabPrefixAutomaton:
+    def __init__(self, vocab: Sequence[str]) -> None: ...
+    def get_order(self) -> Sequence[int]: ...
+    @property
+    def vocab_size(self) -> int: ...
+
 class ReorderedTokenId: ...
 class SearchTree: ...

python/tests/test_vocab.py

@@ -0,0 +1,15 @@
+from mtc_token_healing import VocabPrefixAutomaton
+
+
+def test_vocab_simple():
+    vocab = ["bcd", "abc", "cc", "hello", "world", " ", "yes", "no", "."]
+    order = [5, 8, 1, 0, 2, 3, 7, 4, 6]
+
+    assert len(vocab) == len(order)
+
+    automaton = VocabPrefixAutomaton(vocab)
+
+    assert automaton.vocab_size == len(vocab)
+    assert automaton.get_order() == order
+
+    assert all(vocab[order[i]] < vocab[order[i + 1]] for i in range(len(order) - 1))

src/vocab.rs

@@ -99,7 +99,7 @@ mod _pyo3 {
             Self(Arc::new(VocabPrefixAutomaton::new(vocab)))
         }
 
-        #[pyo3(name = "vocab_size")]
+        #[getter("vocab_size")]
         fn vocab_size_py(&self) -> usize {
             self.vocab.len()
         }