feat: allow any object as extra values

ChieloNewctle · ChieloNewctle · commit 9b059fa74c7c · 2025-08-01T10:06:46.000+08:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,6 +11,9 @@ homepage = "https://github.com/ModelTC/mtc-token-healing"
 documentation = "https://docs.rs/mtc-token-healing"
 authors = ["Chielo Newctle <ChieloNewctle@gmail.com>"]
 
+[workspace.dependencies]
+general-sam = { version = "1.0.1", features = ["trie"] }
+
 [package]
 name = "mtc-token-healing"
 version.workspace = true
@@ -26,7 +29,7 @@ exclude = ["release-plz.toml", ".github", "python"]
 
 [dependencies]
 derive_more = { version = "2.0.1", features = ["deref", "as_ref"] }
-general-sam = { version = "1.0.1", features = ["trie"] }
+general-sam = { workspace = true }
 pyo3 = { version = "0.25.0", optional = true }
 thiserror = "2.0.12"
 tinyvec = { version = "1.9.0", features = ["alloc"] }
diff --git a/python/Cargo.toml b/python/Cargo.toml
@@ -16,4 +16,5 @@ crate-type = ["cdylib"]
 
 [dependencies]
 mtc-token-healing = { version = "0.2.2-dev", path = "..", features = ["pyo3"] }
+general-sam = { workspace = true }
 pyo3 = { version = "0.25.0", features = ["extension-module", "generate-import-lib", "abi3-py310"] }
diff --git a/python/mtc_token_healing/mtc_token_healing.pyi b/python/mtc_token_healing/mtc_token_healing.pyi
@@ -1,4 +1,4 @@
-from typing import Optional, Sequence, Tuple, overload
+from typing import Generic, Optional, Sequence, Tuple, TypeVar, overload
 
 TokenId = int
 SortedTokenId = int
@@ -35,14 +35,16 @@ class VocabPrefixAutomaton:
         self, token_ids: Sequence[TokenId]
     ) -> Sequence[SortedTokenId]: ...
 
-class TokenSeqTrieNode:
+Value = TypeVar("Value")
+
+class TokenSeqTrieNode(Generic[Value]):
     token: int
-    pred_range: Optional[SortedTokenRange]
     parent: int
     subtree_lower: int
     subtree_upper: int
+    depth: int
+    value: Optional[Value]
 
 def dfs_token_seq_trie(
-    token_ids_seq: Sequence[Sequence[int]],
-    pred_rank_ranges: Sequence[SortedTokenRange],
-) -> Tuple[Sequence[TokenSeqTrieNode], int]: ...
+    token_ids_seq_and_values: Sequence[Tuple[Sequence[TokenId], Value]],
+) -> Tuple[Sequence[TokenSeqTrieNode[Value]], int]: ...
diff --git a/python/src/lib.rs b/python/src/lib.rs
@@ -1,39 +1,9 @@
-use std::borrow::Cow;
+mod prefix_dfs;
 
-use ::mtc_token_healing::{
-    dfs_token_seq_trie, SortedTokenRange, TokenId, TokenSeqInput, TokenSeqTrieNode,
-    VocabPrefixAutomaton,
-};
+use ::mtc_token_healing::{SortedTokenRange, TokenId, 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>, usize) {
-    let inputs = token_ids
-        .into_iter()
-        .zip(pred_rank_ranges)
-        .map(|(s, r)| TokenSeqInput {
-            tokens: Cow::Owned(s),
-            pred_range: r,
-        })
-        .collect();
-    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)
-}
+use crate::prefix_dfs::{dfs_token_seq_trie_py, TokenSeqTrieNode};
 
 #[pymodule]
 fn mtc_token_healing(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
diff --git a/python/src/prefix_dfs.rs b/python/src/prefix_dfs.rs
@@ -1,55 +1,53 @@
 use std::{borrow::Cow, convert::Infallible};
 
 use general_sam::{BTreeTransTable, TravelEvent, Trie, TrieNodeAlike};
+use pyo3::{pyclass, pyfunction, pymethods, PyObject, PyResult, Python};
 
-use crate::{SortedTokenRange, TokenId};
+use crate::TokenId;
 
-#[derive(Clone, Debug)]
-#[cfg_attr(feature = "pyo3", ::pyo3::pyclass(get_all, set_all))]
+#[derive(Debug)]
+#[pyclass(get_all, set_all)]
 pub struct TokenSeqTrieNode {
     pub parent: usize,
     pub subtree_lower: usize,
     pub subtree_upper: usize,
+    pub depth: usize,
     pub token: TokenId,
-    pub pred_range: Option<SortedTokenRange>,
+    pub value: Option<PyObject>,
 }
 
-#[cfg(feature = "pyo3")]
-mod _pyo3 {
-    use pyo3::pymethods;
-
-    use super::TokenSeqTrieNode;
-
-    #[pymethods]
-    impl TokenSeqTrieNode {
-        fn __repr__(&self) -> String {
-            let Self {
-                parent,
-                subtree_lower,
-                subtree_upper,
-                token,
-                pred_range,
-            } = self;
-            let pred_range = pred_range
-                .as_ref()
-                .map(|r| r.repr_py())
-                .unwrap_or("None".to_owned());
-            format!(
-                "TokenSeqTrieNode(\
+#[pymethods]
+impl TokenSeqTrieNode {
+    fn __repr__<'py>(&self, py: Python<'py>) -> PyResult<String> {
+        let Self {
+            parent,
+            subtree_lower,
+            subtree_upper,
+            depth,
+            token,
+            value,
+        } = self;
+        let value = value
+            .as_ref()
+            .map(|v| v.call_method0(py, "__repr__")?.extract::<String>(py))
+            .transpose()?
+            .unwrap_or("None".to_owned());
+        Ok(format!(
+            "TokenSeqTrieNode(\
                 token={token}, \
-                pred_range={pred_range}, \
                 parent={parent}, \
                 subtree_lower={subtree_lower}, \
-                subtree_upper={subtree_upper})",
-            )
-        }
+                subtree_upper={subtree_upper}, \
+                depth={depth}, \
+                value={value})",
+        ))
     }
 }
 
 #[derive(Debug)]
 pub struct TokenSeqInput<'a> {
     pub tokens: Cow<'a, [TokenId]>,
-    pub pred_range: SortedTokenRange,
+    pub value: Option<PyObject>,
 }
 
 pub fn dfs_token_seq_trie(inputs: Vec<TokenSeqInput>) -> Vec<TokenSeqTrieNode> {
@@ -80,8 +78,9 @@ pub fn dfs_token_seq_trie(inputs: Vec<TokenSeqInput>) -> Vec<TokenSeqTrieNode> {
                         parent,
                         subtree_lower: dfs_order_id,
                         subtree_upper: dfs_order_id,
+                        depth: 0,
                         token,
-                        pred_range: None,
+                        value: None,
                     });
                 }
                 TravelEvent::Pop(node, _) => {
@@ -99,22 +98,71 @@ pub fn dfs_token_seq_trie(inputs: Vec<TokenSeqInput>) -> Vec<TokenSeqTrieNode> {
         Err(e) => match e {},
     }
 
+    for i in 0..dfs_order.len() {
+        let parent = dfs_order[i].parent;
+        if parent == i {
+            continue;
+        }
+        dfs_order[i].depth = dfs_order[parent].depth + 1;
+    }
+
     for (input, node_id) in inputs.into_iter().zip(seq_last_trie_node_ids) {
         if let Some(id) = rank[node_id] {
-            dfs_order[id].pred_range = Some(input.pred_range);
+            dfs_order[id].value = input.value;
         }
     }
 
     #[cfg(debug_assertions)]
     for (i, node) in dfs_order.iter().enumerate() {
         debug_assert!(node.subtree_lower <= node.subtree_upper);
+        debug_assert!(node.subtree_lower == i);
+        debug_assert!(node.parent <= i);
         if node.parent < node.subtree_lower {
             debug_assert!(node.parent != i);
             let parent = &dfs_order[node.parent];
             debug_assert!(parent.subtree_lower < node.subtree_lower);
             debug_assert!(parent.subtree_upper >= node.subtree_lower);
+        } else {
+            debug_assert!(node.parent == i);
         }
     }
 
     dfs_order
 }
+
+#[pyfunction(name = "dfs_token_seq_trie")]
+pub fn dfs_token_seq_trie_py<'py>(
+    py: Python<'py>,
+    inputs: Vec<(Vec<TokenId>, Option<PyObject>)>,
+) -> (Vec<TokenSeqTrieNode>, usize) {
+    debug_assert!(inputs
+        .iter()
+        .all(|(_, o)| o.as_ref().is_none_or(|v| !v.is_none(py))));
+
+    py.allow_threads(|| {
+        let inputs = inputs
+            .into_iter()
+            .map(|(s, v)| TokenSeqInput {
+                tokens: Cow::Owned(s),
+                value: v,
+            })
+            .collect();
+
+        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.value.is_none() {
+                    res += 1;
+                    continue;
+                }
+                break;
+            }
+            res
+        };
+
+        (nodes, parent_chain_len)
+    })
+}
diff --git a/python/tests/test_dfs_token_seq_trie.py b/python/tests/test_dfs_token_seq_trie.py
@@ -16,29 +16,58 @@ def test_dfs_token_seq_trie():
         SortedTokenRange(1, 10),
         SortedTokenRange(5, 9),
     ]
+    return run_test_dfs_token_seq_trie(tokens_seq, pred_ranges)
 
-    nodes = dfs_token_seq_trie(tokens_seq, pred_ranges)
 
+def test_dfs_trie_value_on_prefix_chain():
+    tokens_seq = [
+        [3, 9],
+        [3, 9, 1, 10, 9, 6, 7],
+        [3, 9, 1, 10, 9, 5],
+        [3, 9, 1, 10, 2],
+        [3, 9, 1, 10],
+        [3, 9, 1, 11],
+    ]
+    pred_ranges = [
+        SortedTokenRange(0, 12),
+        SortedTokenRange(4, 6),
+        SortedTokenRange(4, 7),
+        SortedTokenRange(3, 9),
+        SortedTokenRange(1, 10),
+        SortedTokenRange(5, 9),
+    ]
+    return run_test_dfs_token_seq_trie(tokens_seq, pred_ranges)
+
+
+def run_test_dfs_token_seq_trie(tokens_seq, pred_ranges):
+    nodes, pre_len = dfs_token_seq_trie(list(zip(tokens_seq, pred_ranges)))
+
+    print(f"{pre_len=}")
     print([node.token for node in nodes])
+    print([node.depth for node in nodes])
+    print([node.subtree_upper for node in nodes])
 
-    for i, node in enumerate(nodes):
-        if node.pred_range is None:
+    for q, node in enumerate(nodes):
+        if node.value is None:
             continue
         seq = []
-        for j in range(i):
+        for j in range(q):
             if nodes[j].subtree_upper >= node.subtree_upper:
                 seq.append(nodes[j].token)
         seq.append(node.token)
-        print(seq)
+        print(seq, node.value)
         assert seq in tokens_seq
+        assert pred_ranges[tokens_seq.index(seq)] == node.value
+        assert node.depth + 1 == len(seq)
 
-    for i in range(len(nodes)):
+    for q in range(len(nodes)):
         masks = [
-            j <= i and nodes[j].subtree_upper >= nodes[i].subtree_upper
-            for j in range(len(nodes))
+            k <= q and nodes[k].subtree_upper >= nodes[q].subtree_upper
+            for k in range(len(nodes))
         ]
         print("".join(map(str, map(int, masks))))
 
 
 if __name__ == "__main__":
     test_dfs_token_seq_trie()
+    test_dfs_trie_value_on_prefix_chain()
diff --git a/src/lib.rs b/src/lib.rs
@@ -11,12 +11,10 @@
 //! is the same as walking to the state on the suffix automaton
 //! and gathering information among the subtree of the link tree.
 mod automaton;
-mod prefix_dfs;
 mod token;
 
 pub use crate::{
     automaton::VocabPrefixAutomaton,
-    prefix_dfs::{TokenSeqInput, TokenSeqTrieNode, dfs_token_seq_trie},
     token::{SmallToken, SortedTokenId, SortedTokenRange, TokenId},
 };
 
