feat: better interface

Author: ChieloNewctle

python/mtc_token_healing/mtc_token_healing.pyi

@@ -1,4 +1,4 @@
-from typing import Optional, Sequence, Tuple
+from typing import Optional, Sequence, Tuple, overload
 
 TokenId = int
 SortedTokenId = int
@@ -19,15 +19,30 @@ class VocabPrefixAutomaton:
     def parse_tokens(
         self, token_ids: Sequence[TokenId]
     ) -> Sequence[Tuple[bytes, SortedTokenRange]]: ...
+    def parse_tokens_str_suffix(
+        self, token_ids: Sequence[TokenId]
+    ) -> Sequence[Tuple[str, SortedTokenRange]]: ...
+    @overload
+    def get_original_token_ids(self, sorted_token_id: SortedTokenId) -> TokenId: ...
+    @overload
+    def get_original_token_ids(
+        self, sorted_token_ids: Sequence[SortedTokenId]
+    ) -> Sequence[TokenId]: ...
+    @overload
+    def get_sorted_token_ids(self, token_id: TokenId) -> SortedTokenId: ...
+    @overload
+    def get_sorted_token_ids(
+        self, token_ids: Sequence[TokenId]
+    ) -> Sequence[SortedTokenId]: ...
 
 class TokenSeqTrieNode:
-    token: TokenId
+    token: int
     pred_range: Optional[SortedTokenRange]
     parent: int
     subtree_lower: int
     subtree_upper: int
 
 def dfs_token_seq_trie(
-    token_ids_seq: Sequence[Sequence[TokenId]],
+    token_ids_seq: Sequence[Sequence[int]],
     pred_rank_ranges: Sequence[SortedTokenRange],
-) -> Sequence[TokenSeqTrieNode]: ...
+) -> Tuple[Sequence[TokenSeqTrieNode], int]: ...

python/src/lib.rs

@@ -1,16 +1,16 @@
 use std::borrow::Cow;
 
 use ::mtc_token_healing::{
-    SortedTokenRange, TokenId, TokenSeqInput, TokenSeqTrieNode, VocabPrefixAutomaton,
-    dfs_token_seq_trie,
+    dfs_token_seq_trie, SortedTokenRange, TokenId, TokenSeqInput, TokenSeqTrieNode,
+    VocabPrefixAutomaton,
 };
 use pyo3::prelude::*;
 
 #[pyfunction(name = "dfs_token_seq_trie")]
 fn dfs_token_seq_trie_py(
     token_ids: Vec<Vec<TokenId>>,
     pred_rank_ranges: Vec<SortedTokenRange>,
-) -> Vec<TokenSeqTrieNode> {
+) -> (Vec<TokenSeqTrieNode>, usize) {
     let inputs = token_ids
         .into_iter()
         .zip(pred_rank_ranges)
@@ -19,7 +19,20 @@ fn dfs_token_seq_trie_py(
             pred_range: r,
         })
         .collect();
-    dfs_token_seq_trie(inputs)
+    let nodes = dfs_token_seq_trie(inputs);
+    let parent_chain_len = {
+        let mut res = 0;
+        while res < nodes.len() {
+            let node = &nodes[res];
+            if node.parent == res.saturating_sub(1) && node.pred_range.is_none() {
+                res += 1;
+                continue;
+            }
+            break;
+        }
+        res
+    };
+    (nodes, parent_chain_len)
 }
 
 #[pymodule]

python/tests/test_dfs_token_seq_trie.py

@@ -34,7 +34,7 @@ def test_dfs_token_seq_trie():
 
     for i in range(len(nodes)):
         masks = [
-            j < i and nodes[j].subtree_upper >= nodes[i].subtree_upper
+            j <= i and nodes[j].subtree_upper >= nodes[i].subtree_upper
             for j in range(len(nodes))
         ]
         print("".join(map(str, map(int, masks))))

src/automaton.rs

@@ -120,12 +120,32 @@ impl VocabPrefixAutomaton {
 
 #[cfg(feature = "pyo3")]
 pub mod pyo3 {
-    use pyo3::{Bound, Python, pymethods, types::PyBytes};
+    use pyo3::{Bound, FromPyObject, IntoPyObject, Python, pymethods, types::PyBytes};
 
     use crate::{SortedTokenId, SortedTokenRange, TokenId};
 
     use super::VocabPrefixAutomaton;
 
+    #[derive(Debug, FromPyObject, IntoPyObject)]
+    enum TokenIdSeq {
+        TokenId(TokenId),
+        Seq(Vec<TokenId>),
+    }
+
+    impl TokenIdSeq {
+        fn map<F: FnMut(TokenId) -> TokenId>(self, mut f: F) -> Self {
+            match self {
+                Self::TokenId(id) => Self::TokenId(f(id)),
+                Self::Seq(mut items) => {
+                    items.iter_mut().for_each(|item| {
+                        *item = f(*item);
+                    });
+                    Self::Seq(items)
+                }
+            }
+        }
+    }
+
     #[pymethods]
     impl VocabPrefixAutomaton {
         #[new]
@@ -151,10 +171,11 @@ pub mod pyo3 {
         #[pyo3(name = "parse_bytes")]
         fn parse_bytes_py(
             &self,
+            py: Python<'_>,
             bytes: &[u8],
             start_from: usize,
         ) -> Vec<(usize, SortedTokenRange)> {
-            self.parse_bytes(bytes, start_from)
+            py.allow_threads(|| self.parse_bytes(bytes, start_from))
         }
 
         #[pyo3(name = "parse_tokens")]
@@ -163,10 +184,32 @@ pub mod pyo3 {
             py: Python<'py>,
             tokens: Vec<usize>,
         ) -> Vec<(Bound<'py, PyBytes>, SortedTokenRange)> {
-            self.parse_rev_token_id_seq(tokens.into_iter().rev())
-                .into_iter()
+            let res = py.allow_threads(|| self.parse_rev_token_id_seq(tokens.into_iter().rev()));
+            res.into_iter()
                 .map(|(b, c)| (PyBytes::new(py, &b), c))
                 .collect()
         }
+
+        #[pyo3(name = "parse_tokens_str_suffix")]
+        fn parse_tokens_str_suffix_py(
+            &self,
+            py: Python<'_>,
+            tokens: Vec<usize>,
+        ) -> Vec<(String, SortedTokenRange)> {
+            py.allow_threads(|| {
+                self.parse_rev_token_id_seq(tokens.into_iter().rev())
+                    .into_iter()
+                    .filter_map(|(b, c)| String::from_utf8(b.into()).ok().map(|s| (s, c)))
+                    .collect()
+            })
+        }
+
+        fn get_original_token_ids(&self, py: Python<'_>, seq: TokenIdSeq) -> TokenIdSeq {
+            py.allow_threads(|| seq.map(|id| self.order.get(id as usize).copied().unwrap_or(id)))
+        }
+
+        fn get_sorted_token_ids(&self, py: Python<'_>, seq: TokenIdSeq) -> TokenIdSeq {
+            py.allow_threads(|| seq.map(|id| self.rank.get(id as usize).copied().unwrap_or(id)))
+        }
     }
 }

src/token.rs

@@ -29,7 +29,7 @@ mod _pyo3 {
     impl SortedTokenRange {
         pub(crate) fn repr_py(&self) -> String {
             let Self { lower, upper } = self;
-            format!("SortedTokenRange(lower={}, upper={})", lower, upper)
+            format!("SortedTokenRange(lower={lower}, upper={upper})")
         }
     }