mii/tokenizers
1
0
Fork 0
mirror of https://github.com/mii443/tokenizers.git synced 2025-08-22 16:25:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

More cache options. (#1675)

Browse Source 
* More cache options.

* Fixing error messages.
This commit is contained in:
Nicolas Patry
2024-11-06 18:12:09 +08:00
committed by GitHub
parent 1740bff7a6
commit c6b5c3eab7
5 changed files with 89 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
bindings/python/src/models.rs
View File

@ -534,6 +534,30 @@ impl PyBPE {
)?,
)
}
/// Clears the internal cache
#[pyo3(signature = ())]
#[pyo3(text_signature = "(self)")]
fn _clear_cache(self_: PyRef<Self>) -> PyResult<()> {
let super_ = self_.as_ref();
let mut model = super_.model.write().map_err(|e| {
exceptions::PyException::new_err(format!("Error while clearing BPE cache: {}", e))
})?;
model.clear_cache();
Ok(())
}
/// Resize the internal cache
#[pyo3(signature = (capacity))]
#[pyo3(text_signature = "(self, capacity)")]
fn _resize_cache(self_: PyRef<Self>, capacity: usize) -> PyResult<()> {
let super_ = self_.as_ref();
let mut model = super_.model.write().map_err(|e| {
exceptions::PyException::new_err(format!("Error while resizing BPE cache: {}", e))
})?;
model.resize_cache(capacity);
Ok(())
}
}
/// An implementation of the WordPiece algorithm
@ -858,6 +882,30 @@ impl PyUnigram {
)),
}
}
/// Clears the internal cache
#[pyo3(signature = ())]
#[pyo3(text_signature = "(self)")]
fn _clear_cache(self_: PyRef<Self>) -> PyResult<()> {
let super_ = self_.as_ref();
let mut model = super_.model.write().map_err(|e| {
exceptions::PyException::new_err(format!("Error while clearing Unigram cache: {}", e))
})?;
model.clear_cache();
Ok(())
}
/// Resize the internal cache
#[pyo3(signature = (capacity))]
#[pyo3(text_signature = "(self, capacity)")]
fn _resize_cache(self_: PyRef<Self>, capacity: usize) -> PyResult<()> {
let super_ = self_.as_ref();
let mut model = super_.model.write().map_err(|e| {
exceptions::PyException::new_err(format!("Error while resizing Unigram cache: {}", e))
})?;
model.resize_cache(capacity);
Ok(())
}
}
/// Models Module