scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 145

refactor(rust): simplify tokenizer bridge features

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-20 21:25:19 +00:00
parent 9b2b3d0b90
commit a619b765fc
2 changed files with 17 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -10,13 +10,13 @@ name = "_native"
crate-type = ["cdylib"]
[features]
default = ["abi3", "token-counter-huggingface", "token-counter-tiktoken"]
default = ["abi3", "huggingface", "tiktoken"]
abi3 = ["pyo3/abi3-py310"]
extension-module = ["pyo3/extension-module"]
panic-test = []
token-counter-fast = ["litellm-token-counter/fast"]
token-counter-huggingface = ["litellm-token-counter/huggingface"]
token-counter-tiktoken = ["litellm-token-counter/tiktoken"]
fast = ["litellm-token-counter/fast"]
huggingface = ["litellm-token-counter/huggingface"]
tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
bytes.workspace = true

42
litellm-rust/crates/python-bridge/src/token_counter.rs
View file

@ -1,17 +1,9 @@
use std::sync::Arc;
#[cfg(any(
feature = "token-counter-fast",
feature = "token-counter-huggingface",
feature = "token-counter-tiktoken"
))]
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
use std::{num::NonZero, thread::available_parallelism};
#[cfg(any(
feature = "token-counter-fast",
feature = "token-counter-huggingface",
feature = "token-counter-tiktoken"
))]
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
use litellm_host_python::release_gil;
use litellm_host_python::run_async;
use litellm_token_counter::{
@ -41,45 +33,45 @@ pub(crate) struct TokenCounter {
impl TokenCounter {
#[new]
fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
#[cfg(feature = "token-counter-huggingface")]
#[cfg(feature = "huggingface")]
{
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
}
#[cfg(not(feature = "token-counter-huggingface"))]
#[cfg(not(feature = "huggingface"))]
{
let _ = (py, tokenizer_json);
Err(RustBridgeDeclined::new_err(
"tokenizer backend requires the token-counter-huggingface feature",
"tokenizer backend requires the huggingface feature",
))
}
}
#[staticmethod]
fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
#[cfg(feature = "token-counter-fast")]
#[cfg(feature = "fast")]
{
Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
}
#[cfg(not(feature = "token-counter-fast"))]
#[cfg(not(feature = "fast"))]
{
let _ = (py, tokenizer_json);
Err(RustBridgeDeclined::new_err(
"tokenizer backend requires the token-counter-fast feature",
"tokenizer backend requires the fast feature",
))
}
}
#[staticmethod]
fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
#[cfg(feature = "token-counter-tiktoken")]
#[cfg(feature = "tiktoken")]
{
Self::load(py, || CoreTokenCounter::from_tiktoken(encoding))
}
#[cfg(not(feature = "token-counter-tiktoken"))]
#[cfg(not(feature = "tiktoken"))]
{
let _ = (py, encoding);
Err(RustBridgeDeclined::new_err(
"tokenizer backend requires the token-counter-tiktoken feature",
"tokenizer backend requires the tiktoken feature",
))
}
}
@ -105,11 +97,7 @@ impl TokenCounter {
}
impl TokenCounter {
#[cfg(any(
feature = "token-counter-fast",
feature = "token-counter-huggingface",
feature = "token-counter-tiktoken"
))]
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
fn load(
py: Python<'_>,
load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
@ -122,11 +110,7 @@ impl TokenCounter {
}
}
#[cfg(any(
feature = "token-counter-fast",
feature = "token-counter-huggingface",
feature = "token-counter-tiktoken"
))]
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
fn encode_parallelism() -> usize {
available_parallelism().map_or(1, NonZero::get)
}