propagate to unigram

huggingface · Oct 21, 2024 · ab3236f · ab3236f
1 parent 6e0175a
commit ab3236f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/bindings/python/src/models.rs b/bindings/python/src/models.rs
@@ -537,8 +537,8 @@ impl PyBPE {
 
     fn clear_cache(self_: PyRef<Self>){
         let super_ = self_.as_ref();
-        let model = super_.model.read().unwrap();
-        if let ModelWrapper::BPE(ref mo) = *model {
+        let mut model = super_.model.write().unwrap();
+        if let ModelWrapper::Unigram(ref mut mo) = *model {
             mo.clear_cache()
         } else {
             unreachable!()
@@ -868,6 +868,16 @@ impl PyUnigram {
             )),
         }
     }
+
+    fn clear_cache(self_: PyRef<Self>){
+        let super_ = self_.as_ref();
+        let mut model = super_.model.write().unwrap();
+        if let ModelWrapper::Unigram(ref mut mo) = *model {
+            mo.clear_cache()
+        } else {
+            unreachable!()
+        };
+    }
 }
 
 /// Models Module

diff --git a/tokenizers/src/models/unigram/model.rs b/tokenizers/src/models/unigram/model.rs
@@ -377,6 +377,10 @@ impl Unigram {
         let string = read_to_string(path)?;
         Ok(serde_json::from_str(&string)?)
     }
+
+    pub fn clear_cache(&mut self){
+        self.cache.clear();
+    }
 }
 
 /// Iterator to iterate of vocabulary of the model, and their relative score.